pypcapkit 1.3.3.post1__cp313-none-any.whl

pcapkit/foundation/traceflow/traceflow.py

@@ -0,0 +1,390 @@
+# -*- coding: utf-8 -*-
+# mypy: disable-error-code=dict-item
+"""Base Class
+================
+
+.. module:: pcapkit.foundation.traceflow.traceflow
+
+:mod:`pcapkit.foundation.traceflow.traceflow` contains
+:class:`~pcapkit.foundation.traceflow.traceflow.TraceFlow` only,
+which is an abstract base class for all flow tracing classes.
+
+"""
+import abc
+import collections
+import os
+import sys
+from typing import TYPE_CHECKING, Generic, TypeVar, cast, overload
+
+from dictdumper.dumper import Dumper
+
+from pcapkit.corekit.module import ModuleDescriptor
+from pcapkit.dumpkit.common import make_dumper
+from pcapkit.protocols import __proto__ as protocol_registry
+from pcapkit.protocols.misc.raw import Raw
+from pcapkit.utilities.exceptions import FileExists, RegistryError, stacklevel
+from pcapkit.utilities.warnings import FileWarning, FormatWarning, RegistryWarning, warn
+
+__all__ = ['TraceFlow']
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, DefaultDict, Optional, Type
+
+    from typing_extensions import Literal, Self
+
+    from pcapkit.corekit.infoclass import Info
+    from pcapkit.protocols.protocol import ProtocolBase as Protocol
+
+    CallbackFn = Callable[['_IT'], None]
+
+_DT = TypeVar('_DT')
+_BT = TypeVar('_BT', bound='Info')
+_IT = TypeVar('_IT', bound='Info')
+_PT = TypeVar('_PT', bound='Info')
+
+
+class TraceFlowMeta(abc.ABCMeta):
+    """Meta class to add dynamic support to :class:`TraceFlow`.
+
+    This meta class is used to generate necessary attributes for the
+    :class:`TraceFlow` class. It can be useful to reduce unnecessary
+    registry calls and simplify the customisation process.
+
+    """
+    if TYPE_CHECKING:
+        #: Protocol name of current object.
+        __protocol_name__: 'str'
+        #: Protocol of current object.
+        __protocol_type__: 'Type[Protocol]'
+
+    @property
+    def name(cls) -> 'str':
+        """Protocol name of current object."""
+        if hasattr(cls, '__protocol_name__'):
+            return cls.__protocol_name__
+        return cls.__name__
+
+    @property
+    def protocol(cls) -> 'Type[Protocol]':
+        """Protocol of current object."""
+        if hasattr(cls, '__protocol_type__'):
+            return cls.__protocol_type__
+        return protocol_registry.get(cls.name.upper(), Raw)
+
+
+class TraceFlowBase(Generic[_DT, _BT, _IT, _PT], metaclass=TraceFlowMeta):
+    """Base flow tracing class.
+
+    Arguments:
+        fout: output path
+        format: output format
+        byteorder: output file byte order
+        nanosecond: output nanosecond-resolution file flag
+
+    Note:
+        This class is for internal use only. For customisation, please use
+        :class:`TraceFlow` instead.
+
+    """
+    if TYPE_CHECKING:
+        #: Protocol name of current reassembly object.
+        __protocol_name__: 'str'
+        #: Protocol of current reassembly object.
+        __protocol_type__: 'Type[Protocol]'
+
+        #: List of callback functions upon reassembled datagram.
+        __callback_fn__: 'list[CallbackFn]'
+
+    # Internal data storage for cached properties.
+    __cached__: 'dict[str, Any]'
+
+    ##########################################################################
+    # Defaults.
+    ##########################################################################
+
+    #: DefaultDict[str, tuple[ModuleDescriptor[Dumper] | Type[Dumper], str | None]]:
+    #: Format dumper mapping for writing output files. The values should be a
+    #: tuple representing the module name and class name, or a
+    #: :class:`dictdumper.dumper.Dumper` subclass, and corresponding file extension.
+    __output__ = collections.defaultdict(
+        lambda: (ModuleDescriptor('pcapkit.dumpkit', 'NotImplementedIO'), None),
+        {
+            'pcap': (ModuleDescriptor('pcapkit.dumpkit', 'PCAPIO'), '.pcap'),
+            'cap': (ModuleDescriptor('pcapkit.dumpkit', 'PCAPIO'), '.pcap'),
+            'plist': (ModuleDescriptor('dictdumper', 'PLIST'), '.plist'),
+            'xml': (ModuleDescriptor('dictdumper', 'PLIST'), '.plist'),
+            'json': (ModuleDescriptor('dictdumper', 'JSON'), '.json'),
+            'tree': (ModuleDescriptor('dictdumper', 'Tree'), '.txt'),
+            'text': (ModuleDescriptor('dictdumper', 'Text'), '.txt'),
+            'txt': (ModuleDescriptor('dictdumper', 'Tree'), '.txt'),
+        },
+    )  # type: DefaultDict[str, tuple[ModuleDescriptor[Dumper] | Type[Dumper], str | None]]
+
+    ##########################################################################
+    # Properties.
+    ##########################################################################
+
+    @property
+    def name(self) -> 'str':
+        """Protocol name of current flow tracing object.
+
+        Note:
+            This property is not available as a class
+            attribute.
+
+        """
+        if hasattr(self, '__protocol_name__'):
+            return self.__protocol_name__
+        return type(self).name  # type: ignore[return-value]
+
+    @property
+    def protocol(self) -> 'Type[Protocol]':
+        """Protocol of current flow tracing object.
+
+        Note:
+            This property is not available as a class
+            attribute.
+
+        """
+        if hasattr(self, '__protocol_type__'):
+            return self.__protocol_type__
+        return type(self).protocol  # type: ignore[return-value]
+
+    @property
+    def index(self) -> 'tuple[_IT, ...]':
+        """_IT table for traced flow."""
+        if self._buffer:
+            return self.submit()
+        return tuple(self._stream)
+
+    ##########################################################################
+    # Methods.
+    ##########################################################################
+
+    @classmethod
+    def register_dumper(cls, format: 'str', dumper: 'ModuleDescriptor[Dumper] | Type[Dumper]', ext: 'str') -> 'None':
+        r"""Register a new dumper class.
+
+        Notes:
+            The full qualified class name of the new dumper class
+            should be as ``{dumper.module}.{dumper.name}``.
+
+        Arguments:
+            format: format name
+            dumper: module descriptor or a :class:`dictdumper.dumper.Dumper` subclass
+            ext: file extension
+
+        """
+        if isinstance(dumper, ModuleDescriptor):
+            dumper = dumper.klass
+        if not issubclass(dumper, Dumper):
+            raise RegistryError(f'dumper must be a Dumper subclass, not {dumper!r}')
+        if format in cls.__output__:
+            warn(f'dumper {format} already registered, overwriting', RegistryWarning)
+        cls.__output__[format] = (dumper, ext)
+
+    @classmethod
+    def register_callback(cls, callback: 'CallbackFn', *, index: 'Optional[int]' = None) -> 'None':
+        """Register callback function.
+
+        Arguments:
+            callback: callback function, which will be called
+                when reassembled datagram is obtained, with the
+                list of reassembled datagrams as its only argument
+            index: index of datagram to be called
+
+        """
+        if index is not None:
+            cls.__callback_fn__.insert(index, callback)
+        else:
+            cls.__callback_fn__.append(callback)
+
+    @classmethod
+    def make_fout(cls, fout: 'str' = './tmp', fmt: 'str' = 'pcap') -> 'tuple[Type[Dumper], str | None]':
+        """Make root path for output.
+
+        Args:
+            fout: root path for output
+            fmt: output format
+
+        Returns:
+            Dumper of specified format and file extension of output file.
+
+        Warns:
+            FormatWarning: If ``fmt`` is not supported.
+            FileWarning: If ``fout`` exists and ``fmt`` is :data:`None`.
+
+        Raises:
+            FileExists: If ``fout`` exists and ``fmt`` is **NOT** :data:`None`.
+
+        """
+        output, ext = cls.__output__[fmt]
+        if ext is None:
+            warn(f'Unsupported output format: {fmt}; disabled file output feature',
+                 FormatWarning, stacklevel=stacklevel())
+        if isinstance(output, ModuleDescriptor):
+            output = output.klass
+
+        try:
+            os.makedirs(fout, exist_ok=True)
+        except FileExistsError as error:
+            if ext is None:
+                warn(error.strerror, FileWarning, stacklevel=stacklevel())
+            else:
+                raise FileExists(*error.args).with_traceback(error.__traceback__)
+        return make_dumper(output), ext
+
+    @abc.abstractmethod
+    def dump(self, packet: '_PT') -> 'None':
+        """Dump frame to output files.
+
+        Arguments:
+            packet: a flow packet (:term:`trace.tcp.packet`)
+
+        """
+
+    @overload
+    def trace(self, packet: '_PT', *, output: 'Literal[True]' = ...) -> 'Dumper': ...
+    @overload
+    def trace(self, packet: '_PT', *, output: 'Literal[False]' = ...) -> 'str': ...
+
+    @abc.abstractmethod
+    def trace(self, packet: '_PT', *, output: 'bool' = False) -> 'Dumper | str':
+        """Trace packets.
+
+        Arguments:
+            packet: a flow packet (:term:`trace.tcp.packet`)
+            output: flag if has formatted dumper
+
+        Returns:
+            If ``output`` is :data:`True`, returns the initiated
+            :class:`~dictdumper.dumper.Dumper` object, which will dump data to
+            the output file named after the flow label; otherwise, returns the
+            flow label itself.
+
+        """
+
+    @abc.abstractmethod
+    def submit(self) -> 'tuple[_IT, ...]':
+        """Submit traced TCP flows.
+
+        Returns:
+            Traced TCP flow (:term:`trace.tcp.index`).
+
+        """
+
+    ##########################################################################
+    # Data models.
+    ##########################################################################
+
+    def __new__(cls, *args: 'Any', **kwargs: 'Any') -> 'Self':  # pylint: disable=unused-argument
+        self = super().__new__(cls)
+
+        # NOTE: Assign this attribute after ``__new__`` to avoid shared memory
+        # reference between instances.
+        self.__cached__ = {}
+
+        return self
+
+    def __init__(self, fout: 'Optional[str]', format: 'Optional[str]',  # pylint: disable=redefined-builtin
+                 byteorder: 'Literal["little", "big"]' = sys.byteorder,
+                 nanosecond: bool = False) -> 'None':
+        """Initialise instance.
+
+        Arguments:
+            fout: output path
+            format: output format
+            byteorder: output file byte order
+            nanosecond: output nanosecond-resolution file flag
+
+        """
+        if fout is None:
+            fout = './tmp'
+        if format is None:
+            format = 'pcap'
+
+        #: str: Output root path.
+        self._fproot = fout
+
+        #: dict[_DT, _BT]: Buffer field (:term:`trace.tcp.buffer`).
+        self._buffer = {}  # type: dict[_DT, _BT]
+        #: list[_IT]: Stream index (:term:`trace.tcp.index`).
+        self._stream = []  # type: list[_IT]
+
+        #: Literal['little', 'big']: Output file byte order.
+        self._endian = byteorder
+        #: bool: Output nanosecond-resolution file flag.
+        self._nnsecd = nanosecond
+
+        # dump I/O object
+        fio, ext = self.make_fout(fout, format)
+        #: Type[Dumper]: Dumper class.
+        self._foutio = fio
+        #: Optional[str]: Output file extension.
+        self._fdpext = ext
+
+    def __call__(self, packet: '_PT') -> 'None':
+        """Dump frame to output files.
+
+        Arguments:
+            packet: a flow packet (:term:`trace.tcp.packet`)
+
+        """
+        # trace frame record
+        self.dump(packet)
+
+    def __init_subclass__(cls) -> 'None':
+        """Initialise subclass.
+
+        This method is to be used for generating necessary attributes
+        for the :class:`TraceFlow` class. It can be useful to reduce
+        unnecessary registry calls and simplify the customisation process.
+
+        """
+        cls.__callback_fn__ = []
+
+
+class TraceFlow(TraceFlowBase[_DT, _BT, _IT, _PT], Generic[_DT, _BT, _IT, _PT]):
+    """Base flow tracing class.
+
+    Example:
+
+        Use keyword argument ``protocol`` to specify the protocol
+        name at class definition:
+
+        .. code-block:: python
+
+           class MyProtocol(TraceFlow, protocol='my_protocol'):
+               ...
+
+    Arguments:
+        fout: output path
+        format: output format
+        byteorder: output file byte order
+        nanosecond: output nanosecond-resolution file flag
+
+    """
+
+    def __init_subclass__(cls, /, protocol: 'Optional[str]' = None, *args: 'Any', **kwargs: 'Any') -> 'None':
+        """Initialise subclass.
+
+        This method is to be used for registering the engine class to
+        :class:`~pcapkit.foundation.extraction.Extractor` class.
+
+        Args:
+            name: Protocol name, default to class name.
+            *args: Arbitrary positional arguments.
+            **kwargs: Arbitrary keyword arguments.
+
+        See Also:
+            For more details, please refer to
+            :meth:`pcapkit.foundation.extraction.Extractor.register_traceflow`.
+
+        """
+        if protocol is None:
+            protocol = cast('str', cls.name)
+
+        from pcapkit.foundation.extraction import Extractor
+        Extractor.register_traceflow(protocol.lower(), cls)
+
+        return super().__init_subclass__()

pcapkit/interface/__init__.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=unused-import
+"""User Interface
+====================
+
+.. module:: pcapkit.interface
+
+:mod:`pcapkit.interface` defines several user-oriented
+interfaces, variables, and etc. These interfaces are
+designed to help and simplify the usage of :mod:`pcapkit`.
+
+"""
+
+from pcapkit.interface.core import (APP, DPKT, INET, JSON, LINK, PCAP, PLIST, RAW, TRANS, TREE,
+                                    PCAPKit, PyShark, Scapy, extract, reassemble, trace)
+
+__all__ = [
+    'extract', 'reassemble', 'trace',                       # interface functions
+    'TREE', 'JSON', 'PLIST', 'PCAP',                        # format macros
+    'LINK', 'INET', 'TRANS', 'APP', 'RAW',                  # layer macros
+    'DPKT', 'Scapy', 'PyShark', 'PCAPKit',                  # engine macros
+]

pcapkit/interface/core.py

@@ -0,0 +1,185 @@
+# -*- coding: utf-8 -*-
+"""Core Interface
+====================
+
+.. module:: pcapkit.interface.core
+
+:mod:`pcapkit.interface.core` defines core user-oriented
+interfaces, variables, and etc., which wraps around the
+foundation classes from :mod:`pcapkit.foundation`.
+
+"""
+import io
+import sys
+from typing import TYPE_CHECKING
+
+from pcapkit.foundation.extraction import Extractor
+from pcapkit.foundation.reassembly.ipv4 import IPv4 as IPv4_Reassembly
+from pcapkit.foundation.reassembly.ipv6 import IPv6 as IPv6_Reassembly
+from pcapkit.foundation.reassembly.tcp import TCP as TCP_Reassembly
+from pcapkit.foundation.traceflow.tcp import TCP as TCP_TraceFlow
+from pcapkit.protocols.protocol import ProtocolBase as Protocol
+from pcapkit.utilities.exceptions import FormatError
+
+if TYPE_CHECKING:
+    from typing import IO, Optional, Type
+
+    from typing_extensions import Literal
+
+    from pcapkit.foundation.extraction import Engines, Formats, Layers, Protocols, VerboseHandler
+    from pcapkit.foundation.reassembly.reassembly import ReassemblyBase as Reassembly
+    from pcapkit.foundation.traceflow.traceflow import TraceFlowBase as TraceFlow
+
+__all__ = [
+    'extract', 'reassemble', 'trace',                       # interface functions
+    'TREE', 'JSON', 'PLIST', 'PCAP',                        # format macros
+    'LINK', 'INET', 'TRANS', 'APP', 'RAW',                  # layer macros
+    'DPKT', 'Scapy', 'PyShark', 'PCAPKit',                  # engine macros
+]
+
+# output file formats
+TREE = 'tree'
+JSON = 'json'
+PLIST = 'plist'
+PCAP = 'pcap'
+
+# layer thresholds
+RAW = 'none'
+LINK = 'link'
+INET = 'internet'
+TRANS = 'transport'
+APP = 'application'
+
+# extraction engines
+DPKT = 'dpkt'
+Scapy = 'scapy'
+PCAPKit = 'default'
+PyShark = 'pyshark'
+
+
+def extract(fin: 'Optional[str | IO[bytes]]' = None, fout: 'Optional[str]' = None, format: 'Optional[Formats]' = None,     # basic settings # pylint: disable=redefined-builtin
+            auto: 'bool' = True, extension: 'bool' = True, store: 'bool' = True,                                           # internal settings # pylint: disable=line-too-long
+            files: 'bool' = False, nofile: 'bool' = False, verbose: 'bool | VerboseHandler' = False,                       # output settings # pylint: disable=line-too-long
+            engine: 'Optional[Engines]' = None, layer: 'Optional[Layers] | Type[Protocol]' = None,                         # extraction settings # pylint: disable=line-too-long
+            protocol: 'Optional[Protocols]' = None,                                                                        # extraction settings # pylint: disable=line-too-long
+            reassembly: 'bool' = False, reasm_strict: 'bool' = True, reasm_store: 'bool' = True,                           # reassembly settings # pylint: disable=line-too-long
+            trace: 'bool' = False, trace_fout: 'Optional[str]' = None, trace_format: 'Optional[Formats]' = None,           # trace settings # pylint: disable=line-too-long
+            trace_byteorder: 'Literal["big", "little"]' = sys.byteorder, trace_nanosecond: 'bool' = False,                 # trace settings # pylint: disable=line-too-long
+            ip: 'bool' = False, ipv4: 'bool' = False, ipv6: 'bool' = False, tcp: 'bool' = False,                           # reassembly/trace settings # pylint: disable=line-too-long
+            buffer_size: 'int' = io.DEFAULT_BUFFER_SIZE, buffer_save: 'bool' = False, buffer_path: 'Optional[str]' = None, # buffer settings # pylint: disable=line-too-long
+            no_eof: 'bool' = False) -> 'Extractor':
+    """Extract a PCAP file.
+
+    Arguments:
+        fin: file name to be read or a binary IO object;
+            if file not exist, raise :exc:`FileNotFound`
+        fout: file name to be written
+        format: file format of output
+
+        auto: if automatically run till EOF
+        extension: if check and append extensions to output file
+        store: if store extracted packet info
+
+        files: if split each frame into different files
+        nofile: if no output file is to be dumped
+        verbose: a :obj:`bool` value or a function takes the :class:`Extractor`
+            instance and current parsed frame (depends on engine selected) as
+            parameters to print verbose output information
+
+        engine: extraction engine to be used
+        layer: extract til which layer
+        protocol: extract til which protocol
+
+        reassembly: if perform reassembly
+        reasm_strict: if set strict flag for reassembly
+        reasm_store: if store reassembled datagrams
+
+        trace: if trace TCP traffic flows
+        trace_fout: path name for flow tracer if necessary
+        trace_format: output file format of flow tracer
+        trace_byteorder: output file byte order
+        trace_nanosecond: output nanosecond-resolution file flag
+
+        ip: if record data for IPv4 & IPv6 reassembly (must be used with ``reassembly=True``)
+        ipv4: if perform IPv4 reassembly (must be used with ``reassembly=True``)
+        ipv6: if perform IPv6 reassembly (must be used with ``reassembly=True``)
+        tcp: if perform TCP reassembly and/or flow tracing
+            (must be used with ``reassembly=True`` or ``trace=True``)
+
+        buffer_size: buffer size for reading input file (for :class:`~pcapkit.corekit.io.SeekableReader` only)
+        buffer_save: if save buffer to file (for :class:`~pcapkit.corekit.io.SeekableReader` only)
+        buffer_path: path name for buffer file if necessary (for :class:`~pcapkit.corekit.io.SeekableReader` only)
+
+        no_eof: if not raise :exc:`EOFError` when reach EOF
+
+    Returns:
+        An :class:`~pcapkit.foundation.extraction.Extractor` object.
+
+    """
+    if isinstance(layer, type) and issubclass(layer, Protocol):
+        layer = (layer.__layer__ or 'none').lower()  # type: ignore[assignment]
+
+    return Extractor(fin=fin, fout=fout, format=format,
+                     store=store, files=files, nofile=nofile,
+                     auto=auto, verbose=verbose, extension=extension,
+                     engine=engine, layer=layer, protocol=protocol,  # type: ignore[arg-type]
+                     ip=ip, ipv4=ipv4, ipv6=ipv6, tcp=tcp,
+                     reassembly=reassembly, reasm_store=reasm_store, reasm_strict=reasm_strict,
+                     trace=trace, trace_fout=trace_fout, trace_format=trace_format,
+                     trace_byteorder=trace_byteorder, trace_nanosecond=trace_nanosecond,
+                     buffer_size=buffer_size, buffer_path=buffer_path, buffer_save=buffer_save,
+                     no_eof=no_eof)
+
+
+def reassemble(protocol: 'str | Type[Protocol]', strict: 'bool' = False) -> 'Reassembly':
+    """Reassemble fragmented datagrams.
+
+    Arguments:
+        protocol: protocol to be reassembled
+        strict: if return all datagrams (including those not implemented) when submit
+
+    Returns:
+        A :class:`~pcapkit.foundation.reassembly.reassembly.Reassembly` object of corresponding protocol.
+
+    Raises:
+        FormatError: If ``protocol`` is **NOT** any of IPv4, IPv6 or TCP.
+
+    """
+    if isinstance(protocol, type) and issubclass(protocol, Protocol):
+        protocol = protocol.id()[0]
+
+    if protocol == 'IPv4':
+        return IPv4_Reassembly(strict=strict)
+    if protocol == 'IPv6':
+        return IPv6_Reassembly(strict=strict)
+    if protocol == 'TCP':
+        return TCP_Reassembly(strict=strict)
+    raise FormatError(f'Unsupported reassembly protocol: {protocol}')
+
+
+def trace(protocol: 'str | Type[Protocol]', fout: 'Optional[str]',
+          format: 'Optional[str]',  # pylint: disable=redefined-builtin
+          byteorder: 'Literal["little", "big"]' = sys.byteorder,
+          nanosecond: bool = False) -> 'TraceFlow':
+    """Trace flows.
+
+    Arguments:
+        protocol: protocol to be reassembled
+        fout: output path
+        format: output format
+        byteorder: output file byte order
+        nanosecond: output nanosecond-resolution file flag
+
+    Returns:
+        A :class:`~pcapkit.foundation.traceflow.traceflow.TraceFlow` object.
+
+    Raises:
+        FormatError: If ``protocol`` is **NOT** TCP.
+
+    """
+    if isinstance(protocol, type) and issubclass(protocol, Protocol):
+        protocol = protocol.id()[0]
+
+    if protocol == 'TCP':
+        return TCP_TraceFlow(fout=fout, format=format, byteorder=byteorder, nanosecond=nanosecond)
+    raise FormatError(f'Unsupported flow tracing protocol: {protocol}')