traffic-taffy 0.5.8__py3-none-any.whl → 0.6.1__py3-none-any.whl

traffic_taffy/__init__.py

@@ -1 +1 @@
-__VERSION__ = "0.5.8"
+__VERSION__ = "0.6.1"

traffic_taffy/compare.py

@@ -1,16 +1,35 @@
+"""The primary statistical packet comparison engine."""
+
+from __future__ import annotations
 from logging import debug, error
-from typing import List
+from typing import List, TYPE_CHECKING
 import datetime as dt
 from datetime import datetime
 
+if TYPE_CHECKING:
+    from argparse import ArgumentParser, Namespace
+
 from traffic_taffy.comparison import Comparison
 from traffic_taffy.dissectmany import PCAPDissectMany
 from traffic_taffy.dissector import PCAPDissectorLevel
 from traffic_taffy.dissection import Dissection
 
+from dataclasses import dataclass
+
+
+@dataclass
+class Report:
+    delta_percentage: float
+    delta_absolute: int
+    total: int
+    left_count: int
+    right_count: int
+    left_percentage: float
+    right_percentage: float
+
 
 class PcapCompare:
-    "Takes a set of PCAPs to then perform various comparisons upon"
+    """Take a set of PCAPs to then perform various comparisons upon."""
 
     REPORT_VERSION: int = 2
 
@@ -25,11 +44,12 @@ class PcapCompare:
         bin_size: int | None = None,
         dissection_level: PCAPDissectorLevel = PCAPDissectorLevel.COUNT_ONLY,
         between_times: List[int] | None = None,
-        ignore_list: List[str] = [],
+        ignore_list: List[str] | None = None,
         layers: List[str] | None = None,
         force_load: bool = False,
         force_overwrite: bool = False,
     ) -> None:
+        """Create a compare object."""
         self.pcap_files = pcap_files
         self.deep = deep
         self.maximum_count = maximum_count
@@ -39,30 +59,31 @@ class PcapCompare:
         self.between_times = between_times
         self.bin_size = bin_size
         self.cache_file_suffix = cache_file_suffix
-        self.ignore_list = ignore_list
+        self.ignore_list = ignore_list or []
         self.layers = layers
         self.force_overwrite = force_overwrite
         self.force_load = force_load
 
     @property
-    def pcap_files(self):
+    def pcap_files(self) -> List[str]:
+        """List of pcap files being compared."""
         return self._pcap_files
 
     @pcap_files.setter
-    def pcap_files(self, new_pcap_files):
+    def pcap_files(self, new_pcap_files: List[str]) -> None:
         self._pcap_files = new_pcap_files
 
     @property
-    def reports(self):
+    def reports(self) -> List[dict]:
+        """List of reports generated by the comparison."""
         return self._reports
 
     @reports.setter
-    def reports(self, newvalue):
+    def reports(self, newvalue: List[dict]) -> None:
         self._reports = newvalue
 
     def compare_dissections(self, left_side: dict, right_side: dict) -> dict:
-        "compares the results from two reports"
-
+        """Compare two dissections."""
         report = {}
 
         keys = set(left_side.keys())
@@ -79,7 +100,7 @@ class PcapCompare:
             right_side_total = sum(right_side[key].values())
 
             new_left_count = 0
-            for subkey in left_side[key].keys():
+            for subkey in left_side[key]:
                 delta_percentage = 0.0
                 total = 0
                 if subkey in right_side[key]:
@@ -99,18 +120,18 @@ class PcapCompare:
                     new_left_count += 1
 
                 delta_absolute = right_count - left_count
-                report[key][subkey] = {
-                    "delta_percentage": delta_percentage,
-                    "delta_absolute": delta_absolute,
-                    "total": total,
-                    "left_count": left_count,
-                    "right_count": right_count,
-                    "left_percentage": left_percentage,
-                    "right_percentage": right_percentage,
-                }
+                report[key][subkey] = Report(
+                    delta_percentage=delta_percentage,
+                    delta_absolute=delta_absolute,
+                    total=total,
+                    left_count=left_count,
+                    right_count=right_count,
+                    left_percentage=left_percentage,
+                    right_percentage=right_percentage,
+                )
 
             new_right_count = 0
-            for subkey in right_side[key].keys():
+            for subkey in right_side[key]:
                 if subkey not in report[key]:
                     delta_percentage = 1.0
                     total = right_side[key][subkey]
@@ -120,15 +141,15 @@ class PcapCompare:
                     right_percentage = right_side[key][subkey] / right_side_total
                     new_right_count += 1  # this value wasn't in the left
 
-                    report[key][subkey] = {
-                        "delta_percentage": delta_percentage,
-                        "delta_absolute": right_count,
-                        "total": total,
-                        "left_count": left_count,
-                        "right_count": right_count,
-                        "left_percentage": left_percentage,
-                        "right_percentage": right_percentage,
-                    }
+                    report[key][subkey] = Report(
+                        delta_percentage=delta_percentage,
+                        delta_absolute=right_count,
+                        total=total,
+                        left_count=left_count,
+                        right_count=right_count,
+                        left_percentage=left_percentage,
+                        right_percentage=right_percentage,
+                    )
 
             if right_side_total == 0:
                 right_percent = 100
@@ -140,19 +161,20 @@ class PcapCompare:
             else:
                 left_percent = new_left_count / left_side_total
 
-            report[key][Dissection.NEW_RIGHT_SUBKEY] = {
-                "delta_absolute": new_right_count - new_left_count,
-                "total": new_left_count + new_right_count,
-                "left_count": new_left_count,
-                "right_count": new_right_count,
-                "left_percentage": left_percent,
-                "right_percentage": right_percent,
-                "delta_percentage": right_percent - left_percent,
-            }
+            report[key][Dissection.NEW_RIGHT_SUBKEY] = Report(
+                delta_absolute=new_right_count - new_left_count,
+                total=new_left_count + new_right_count,
+                left_count=new_left_count,
+                right_count=new_right_count,
+                left_percentage=left_percent,
+                right_percentage=right_percent,
+                delta_percentage=right_percent - left_percent,
+            )
 
         return Comparison(report)
 
     def load_pcaps(self) -> None:
+        """Load all pcaps into memory and dissect them."""
         # load the first as a reference pcap
         pdm = PCAPDissectMany(
             self.pcap_files,
@@ -167,17 +189,16 @@ class PcapCompare:
             force_load=self.force_load,
             force_overwrite=self.force_overwrite,
         )
-        results = pdm.load_all()
-        return results
+        return pdm.load_all()
 
     def compare(self) -> List[Comparison]:
-        "Compares each pcap against the original source"
-
+        """Compare each pcap as requested."""
         dissections = self.load_pcaps()
         self.compare_all(dissections)
         return self.reports
 
-    def compare_all(self, dissections) -> List[Comparison]:
+    def compare_all(self, dissections: List[Dissection]) -> List[Comparison]:
+        """Compare all loaded pcaps."""
         reports = []
         if len(self.pcap_files) > 1:
             # multiple file comparison
@@ -197,9 +218,8 @@ class PcapCompare:
                 error(
                     "the requested pcap data was not long enough to compare against itself"
                 )
-                raise ValueError(
-                    "not enough of a single capture file to time-bin the results"
-                )
+                errorstr: str = "not large enough pcap file"
+                raise ValueError(errorstr)
             debug(
                 f"found {len(timestamps)} timestamps from {timestamps[2]} to {timestamps[-1]}"
             )
@@ -249,7 +269,10 @@ class PcapCompare:
         return reports
 
 
-def compare_add_parseargs(compare_parser, add_subgroup: bool = True):
+def compare_add_parseargs(
+    compare_parser: ArgumentParser, add_subgroup: bool = True
+) -> ArgumentParser:
+    """Add common comparison arguments."""
     if add_subgroup:
         compare_parser = compare_parser.add_argument_group("Comparison result options")
 
@@ -285,17 +308,26 @@ def compare_add_parseargs(compare_parser, add_subgroup: bool = True):
     )
 
     compare_parser.add_argument(
-        "-T",
-        "--between-times",
-        nargs=2,
-        type=int,
-        help="For single files, only display results between these timestamps",
+        "-s",
+        "--sort-by",
+        default="delta%",
+        type=str,
+        help="Sort report entries by this column",
     )
 
+    # compare_parser.add_argument(
+    #     "-T",
+    #     "--between-times",
+    #     nargs=2,
+    #     type=int,
+    #     help="For single files, only display results between these timestamps",
+    # )
+
     return compare_parser
 
 
-def get_comparison_args(args):
+def get_comparison_args(args: Namespace) -> dict:
+    """Return a dict of comparison parameters from arguments."""
     return {
         "maximum_count": args.packet_count or 0,
         "print_threshold": float(args.print_threshold) / 100.0,

traffic_taffy/comparison.py

@@ -1,26 +1,34 @@
-from typing import Dict
+"""A simple data storage module to hold comparison data."""
+
+from __future__ import annotations
+from typing import Dict, Any
 
 
 class Comparison:
+    """A simple data storage class to hold comparison data."""
+
     def __init__(self, contents: list, title: str = ""):
+        """Create a Comparison class from contents."""
         self.contents = contents
         self.title: str = title
-        self.printing_arguments: Dict[str] = {}
+        self.printing_arguments: Dict[str, Any] = {}
 
     # title
     @property
     def title(self) -> str:
+        """The title of this comparison."""
         return self._title
 
     @title.setter
-    def title(self, new_title):
+    def title(self, new_title: str) -> None:
         self._title = new_title
 
     # report contents -- actual data
     @property
-    def contents(self):
+    def contents(self) -> None:
+        """The contents of this comparison."""
         return self._contents
 
     @contents.setter
-    def contents(self, new_contents):
+    def contents(self, new_contents: str) -> None:
         self._contents = new_contents

traffic_taffy/dissection.py

@@ -1,15 +1,16 @@
 """A Dissection class stores the results of a PCAP enumeration."""
 
 from __future__ import annotations
-import os
 from collections import defaultdict, Counter
-from typing import Any
+from typing import Any, Dict, ClassVar
 from logging import debug, info, error, warning
 from enum import Enum
 import msgpack
 import ipaddress
 from typing import List
 from copy import deepcopy
+from pathlib import Path
+from traffic_taffy import __VERSION__ as VERSION
 
 
 class PCAPDissectorLevel(Enum):
@@ -17,6 +18,7 @@ class PCAPDissectorLevel(Enum):
 
     COUNT_ONLY = 1
     THROUGH_IP = 2
+    COMMON_LAYERS = 3
     DETAILED = 10
 
 
@@ -24,13 +26,15 @@ class Dissection:
     """Class to store the data from an enumerated pcap."""
 
     DISSECTION_KEY: str = "PCAP_DISSECTION_VERSION"
-    DISSECTION_VERSION: int = 7
+    DISSECTION_VERSION: int = 8
 
     TOTAL_COUNT: str = "__TOTAL__"
     TOTAL_SUBKEY: str = "packet"
     WIDTH_SUBKEY: str = "__WIDTH__"
     NEW_RIGHT_SUBKEY: str = "__NEW_VALUES__"
 
+    PRINTABLE_LENGTH: int = 40
+
     def __init__(
         self: Dissection,
         pcap_file: str,
@@ -171,11 +175,13 @@ class Dissection:
         """Load the dissection data from a cache."""
         if not self.pcap_file or not isinstance(self.pcap_file, str):
             return None
-        if not os.path.exists(self.pcap_file + self.cache_file_suffix):
+        if not Path(self.pcap_file + self.cache_file_suffix).exists():
             return None
 
         cached_file = self.pcap_file + self.cache_file_suffix
-        cached_contents = self.load_saved(cached_file, dont_overwrite=True)
+        cached_contents = self.load_saved(
+            cached_file, dont_overwrite=True, force_load=force_load
+        )
 
         ok_to_load = True
 
@@ -208,9 +214,10 @@ class Dissection:
                     # loading a more detailed cache is ok
                     continue
 
-                if parameter == "pcap_file" and os.path.basename(
-                    specified
-                ) == os.path.basename(cached):
+                if (
+                    parameter == "pcap_file"
+                    and Path(specified).name == Path(cached).name
+                ):
                     # as long as the basename is ok, we'll assume it's a different path
                     continue
 
@@ -248,13 +255,14 @@ class Dissection:
             self.save(where)
 
     def save(self: Dissection, where: str) -> None:
-        """Saves a generated dissection to a msgpack file."""
+        """Save a generated dissection to a msgpack file."""
         # wrap the report in a version header
         versioned_cache = {
             self.DISSECTION_KEY: self.DISSECTION_VERSION,
             "file": self.pcap_file,
             "parameters": {},
             "dissection": self.data,
+            "created_by": "traffic-taffy " + VERSION,
         }
 
         for parameter in self.parameters:
@@ -303,10 +311,10 @@ class Dissection:
                     ] = versioned_cache["dissection"][timestamp][key][subkey]
                     del versioned_cache["dissection"][timestamp][key][subkey]
 
-        with open(where, "wb") as saveto:
+        with Path(where).open("wb") as saveto:
             msgpack.dump(versioned_cache, saveto)
 
-    def load_saved_contents(self: Dissection, versioned_cache: dict):
+    def load_saved_contents(self: Dissection, versioned_cache: dict) -> None:
         """Set parameters from the cache."""
         # set the local parameters from the cache
         for parameter in self.parameters:
@@ -315,9 +323,14 @@ class Dissection:
         # load the data
         self.data = versioned_cache["dissection"]
 
-    def load_saved(self: Dissection, where: str, dont_overwrite: bool = False) -> dict:
+    def load_saved(
+        self: Dissection,
+        where: str,
+        dont_overwrite: bool = False,
+        force_load: bool = False,
+    ) -> dict:
         """Load a saved report from a cache file."""
-        with open(where, "rb") as cache_file:
+        with Path(where).open("rb") as cache_file:
             contents = msgpack.load(cache_file, strict_map_key=False)
 
         # convert the ignore list to a set (msgpack doesn't do sets)
@@ -326,7 +339,7 @@ class Dissection:
         )
 
         # check that the version header matches something we understand
-        if contents[self.DISSECTION_KEY] != self.DISSECTION_VERSION:
+        if not force_load and contents[self.DISSECTION_KEY] != self.DISSECTION_VERSION:
             raise ValueError(
                 "improper saved dissection version: report version = "
                 + str(contents[self.DISSECTION_KEY])
@@ -413,8 +426,9 @@ class Dissection:
                 value = "0x" + value.hex()
             else:
                 value = "[unprintable]"
-        if len(value) > 40:
-            value = value[0:40] + "..."  # truncate to reasonable
+        if len(value) > Dissection.PRINTABLE_LENGTH:
+            # truncate to reasonable
+            value = value[0 : Dissection.PRINTABLE_LENGTH] + "..."
         return value
 
     @staticmethod
@@ -428,7 +442,7 @@ class Dissection:
         return ":".join(map(two_hex, value))
 
     # has to go at the end to pick up the above function names
-    DISPLAY_TRANSFORMERS: dict = {
+    DISPLAY_TRANSFORMERS: ClassVar[Dict[str, callable]] = {
         "Ethernet.IP.src": ipaddress.ip_address,
         "Ethernet.IP.dst": ipaddress.ip_address,
         "Ethernet.IP6.src": ipaddress.ip_address,

traffic_taffy/dissectmany.py

@@ -1,13 +1,25 @@
-from traffic_taffy.dissector import PCAPDissector
-from pcap_parallel import PCAPParallel
+"""A module for dissecting a number of PCAP files."""
+
+from __future__ import annotations
 from concurrent.futures import ProcessPoolExecutor
 from logging import info
 import copy
 import multiprocessing
+from pcap_parallel import PCAPParallel
+from typing import List, TYPE_CHECKING
+
+from traffic_taffy.dissector import PCAPDissector
+
+if TYPE_CHECKING:
+    from io import BufferedIOBase
+    from traffic_taffy.dissection import Dissection
 
 
 class PCAPDissectMany:
-    def __init__(self, pcap_files, *args, **kwargs):
+    """A class for dissecting a number of PCAP files."""
+
+    def __init__(self, pcap_files: List[str], *args: list, **kwargs: dict):
+        """Create a PCAPDissectMany instance."""
         self.pcap_files = pcap_files
         self.args = args
         self.kwargs = kwargs
@@ -20,7 +32,8 @@ class PCAPDissectMany:
             # Note: this may undercount due to int flooring()
             self.maximum_cores = int(multiprocessing.cpu_count() / len(self.pcap_files))
 
-    def load_pcap_piece(self, pcap_io_buffer):
+    def load_pcap_piece(self, pcap_io_buffer: BufferedIOBase) -> Dissection:
+        """Load one piece of a pcap from a buffer."""
         kwargs = copy.copy(self.kwargs)
         # force false for actually loading
         kwargs["cache_results"] = False
@@ -36,12 +49,10 @@ class PCAPDissectMany:
 
     def load_pcap(
         self,
-        pcap_file,
-        split_size=None,
-        maximum_count: int = 0,
-        force_overwrite: bool = False,
-        force_load: bool = False,
-    ):
+        pcap_file: str,
+        split_size: int | None = None,
+    ) -> Dissection:
+        """Load one pcap file."""
         pd = PCAPDissector(
             pcap_file,
             *self.args,
@@ -55,23 +66,34 @@ class PCAPDissectMany:
             return dissection
 
         info(f"processing {pcap_file}")
-        ps = PCAPParallel(
-            pcap_file,
-            split_size=split_size,
-            callback=self.load_pcap_piece,
-            maximum_count=self.kwargs.get("maximum_count", 0),
-            maximum_cores=self.maximum_cores,
-        )
-        results = ps.split()
+        if isinstance(pcap_file, str) and (
+            pcap_file.endswith(".dnstap") or pcap_file.endswith(".tap")
+        ):
+            # deal with dnstap files
+
+            # the Dissector already handles loading a dnstap engine
+            # TODO(hardaker): see if we can use a splitter here with the framing chunks
+            dissection = pd.load()
+
+        else:  # assume pcap
+            ps = PCAPParallel(
+                pcap_file,
+                split_size=split_size,
+                callback=self.load_pcap_piece,
+                maximum_count=self.kwargs.get("maximum_count", 0),
+                maximum_cores=self.maximum_cores,
+            )
+            results = ps.split()
 
-        # the data is coming back in (likely overlapping) chunks, and
-        # we need to merge them together
-        dissection = results.pop(0).result()
-        dissection.pcap_file = pcap_file  # splitting has the wrong name
-        for result in results:
-            dissection.merge(result.result())
+            # the data is coming back in (likely overlapping) chunks, and
+            # we need to merge them together
+            dissection = results.pop(0).result()
+            dissection.pcap_file = pcap_file  # splitting has the wrong name
+            for result in results:
+                dissection.merge(result.result())
 
-        dissection.calculate_metadata()
+            # recalculate metadata now that merges have happened
+            dissection.calculate_metadata()
 
         if self.kwargs.get("cache_results"):
             # create a dissector just to save the cache
@@ -83,7 +105,10 @@ class PCAPDissectMany:
 
         return dissection
 
-    def load_all(self, return_as_list: bool = False, dont_fork: bool = False):
+    def load_all(
+        self, return_as_list: bool = False, dont_fork: bool = False
+    ) -> List[Dissection]:
+        """Load all PCAPs in parallel."""
         if dont_fork:
             # handle each one individually -- typically for inserting debugging stops
             dissections = []
@@ -96,5 +121,5 @@ class PCAPDissectMany:
         with ProcessPoolExecutor() as executor:
             dissections = executor.map(self.load_pcap, self.pcap_files)
             if return_as_list:  # convert from generator
-                dissections = [x for x in dissections]
+                dissections = list(dissections)
             return dissections

traffic_taffy/dissector.py

@@ -97,7 +97,16 @@ class PCAPDissector:
 
         engine = None
         args = self.dissection_args()
-        if (
+
+        if isinstance(self.pcap_file, str) and (
+            self.pcap_file.endswith(".dnstap") or self.pcap_file.endswith(".tap")
+        ):
+            # we delay loading until the module and its requirements are needed
+            from traffic_taffy.dissector_engine.dnstap import DissectionEngineDNStap
+
+            engine = DissectionEngineDNStap(*args)
+
+        elif (
             self.dissector_level == PCAPDissectorLevel.DETAILED
             or self.dissector_level == PCAPDissectorLevel.DETAILED.value
         ):
@@ -181,7 +190,13 @@ def dissector_add_parseargs(parser, add_subgroup: bool = True):
             "Ethernet.IP.TCP.ack",
             "Ethernet.IPv6.TCP.seq",
             "Ethernet.IPv6.TCP.ack",
+            "Ethernet.IPv6.TCP.Raw.load",
+            "Ethernet.IP.UDP.Raw.load",
             "Ethernet.IP.UDP.DNS.id",
+            "Ethernet.IP.ICMP.IP in ICMP.UDP in ICMP.chksum",
+            "Ethernet.IP.ICMP.IP in ICMP.UDP in ICMP.Raw.load",
+            "Ethernet.IP.ICMP.IP in ICMP.chksum",
+            "Ethernet.IP.ICMP.IP in ICMP.id",
             "Ethernet.IP.TCP.DNS.id",
             "Ethernet.IPv6.UDP.DNS.id",
             "Ethernet.IPv6.TCP.DNS.id",
@@ -189,6 +204,9 @@ def dissector_add_parseargs(parser, add_subgroup: bool = True):
             "Ethernet.IP.chksum",
             "Ethernet.IP.UDP.chksum",
             "Ethernet.IP.TCP.chksum",
+            "Ethernet.IP.TCP.window",
+            "Ethernet.IP.TCP.Raw.load",
+            "Ethernet.IP.UDP.Raw.load",
             "Ethernet.IPv6.UDP.chksum",
             "Ethernet.IPv6.fl",
             "Ethernet.IP.ICMP.chksum",
@@ -197,6 +215,8 @@ def dissector_add_parseargs(parser, add_subgroup: bool = True):
             "Ethernet.IP.TCP.Padding.load",
             "Ethernet.IPv6.TCP.chksum",
             "Ethernet.IPv6.plen",
+            "Ethernet.IP.TCP.Encrypted Content.load",
+            "Ethernet.IP.TCP.TLS.TLS.Raw.load",
         ],
         nargs="*",
         type=str,
@@ -302,6 +322,7 @@ def check_dissector_level(level: int):
     current_dissection_levels = [
         PCAPDissectorLevel.COUNT_ONLY.value,
         PCAPDissectorLevel.THROUGH_IP.value,
+        PCAPDissectorLevel.COMMON_LAYERS.value,
         PCAPDissectorLevel.DETAILED.value,
     ]
     if level not in current_dissection_levels:

traffic_taffy/dissector_engine/__init__.py

@@ -11,7 +11,7 @@ class DissectionEngine:
         pcap_filter: str = "",
         maximum_count: int = 0,
         bin_size: int = 0,
-        dissector_level: PCAPDissectorLevel = PCAPDissectorLevel.DETAILED,
+        dissector_level: PCAPDissectorLevel = PCAPDissectorLevel.COMMON_LAYERS,
         cache_file_suffix: str = "pkl",
         ignore_list: list = [],
         layers: List[str] | None = None,
@@ -25,6 +25,22 @@ class DissectionEngine:
         self.ignore_list = set(ignore_list)
         self.layers = layers
 
+    def start_packet(
+        self, timestamp: int, dissection: Dissection | None = None
+    ) -> None:
+        if not dissection:
+            dissection = self.dissection
+
+        # set and bin-ize the timestamp
+        dissection.timestamp = int(timestamp)
+        if dissection.bin_size:
+            dissection.timestamp = (
+                dissection.timestamp - dissection.timestamp % dissection.bin_size
+            )
+
+        # increment the base counter for all packets
+        dissection.incr(Dissection.TOTAL_COUNT, dissection.TOTAL_SUBKEY)
+
     def init_dissection(self) -> Dissection:
         self.dissection = Dissection(
             pcap_file=self.pcap_file,
@@ -36,3 +52,10 @@ class DissectionEngine:
             ignore_list=self.ignore_list,
         )
         return self.dissection
+
+    def load(self) -> Dissection:
+        """Load the capture file into memory."""
+        self.init_dissection()
+        self.load_data()
+        self.dissection.calculate_metadata()
+        return self.dissection