pypcapkit 1.3.3.post1__cp313-none-any.whl

pcapkit/foundation/traceflow/__init__.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=unused-import
+"""Trace Flows
+=================
+
+.. module:: pcapkit.foundation.traceflow
+
+:mod:`pcapkit.traceflow` implements flow tracing functions for
+:mod:`pcapkit` package.
+
+.. note::
+
+   This was implemented as the demand of my mate
+   `@gousaiyang <https://github.com/gousaiyang>`__.
+
+"""
+# Base Class for TraceFlow
+from pcapkit.foundation.traceflow.traceflow import TraceFlow
+
+# TCP Flow Tracing
+from pcapkit.foundation.traceflow.tcp import TCP as TCP_TraceFlow
+
+__all__ = [
+    'TCP_TraceFlow',
+]
+
+
+from typing import TYPE_CHECKING
+
+from pcapkit.corekit.infoclass import Info, info_final
+
+if TYPE_CHECKING:
+    from typing import Optional
+
+
+@info_final
+class TraceFlowManager(Info):
+    """TraceFlow Manager."""
+
+    #: TCP reassembly.
+    tcp: 'TCP_TraceFlow'
+
+    if TYPE_CHECKING:
+        def __init__(self, tcp: 'Optional[TCP_TraceFlow]') -> 'None': ...  # pylint: disable=unused-argument,super-init-not-called,multiple-statements,line-too-long

pcapkit/foundation/traceflow/data/__init__.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+"""data models for flow tracing"""
+
+# TCP flow tracing
+from pcapkit.foundation.traceflow.data.tcp import Buffer as TCP_Buffer
+from pcapkit.foundation.traceflow.data.tcp import BufferID as TCP_BufferID
+from pcapkit.foundation.traceflow.data.tcp import Index as TCP_Index
+from pcapkit.foundation.traceflow.data.tcp import Packet as TCP_Packet
+
+__all__ = [
+    'TCP_Buffer', 'TCP_BufferID', 'TCP_Index', 'TCP_Packet',
+]
+
+from typing import TYPE_CHECKING
+
+from pcapkit.corekit.infoclass import Info, info_final
+
+if TYPE_CHECKING:
+    from typing import Optional
+
+
+@info_final
+class TraceFlowData(Info):
+    """Data storage for flow tracing."""
+
+    #: TCP traced flows.
+    tcp: 'tuple[TCP_Index, ...]'
+
+    if TYPE_CHECKING:
+        def __init__(self, tcp: 'Optional[tuple[TCP_Index, ...]]') -> 'None': ...  # pylint: disable=unused-argument,super-init-not-called,multiple-statements,line-too-long

pcapkit/foundation/traceflow/data/tcp.py

@@ -0,0 +1,105 @@
+# -*- coding: utf-8 -*-
+"""data models for TCP flow tracing"""
+
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+from pcapkit.corekit.infoclass import Info, info_final
+from pcapkit.utilities.compat import Tuple
+
+__all__ = ['BufferID', 'Packet', 'Buffer', 'Index']
+
+if TYPE_CHECKING:
+    from ipaddress import IPv4Address, IPv6Address
+    from typing import Any, Optional
+
+    from dictdumper.dumper import Dumper
+    from typing_extensions import TypeAlias
+
+    from pcapkit.const.reg.linktype import LinkType as Enum_LinkType
+    from pcapkit.protocols.data.misc.pcap.frame import Frame as Data_Frame
+
+_AT = TypeVar('_AT', 'IPv4Address', 'IPv6Address')
+
+#: Buffer ID.
+BufferID: 'TypeAlias' = Tuple[_AT, int, _AT, int]
+
+
+@info_final
+class Packet(Info, Generic[_AT]):
+    """Data structure for **TCP flow tracing**.
+
+    See Also:
+        * :meth:`pcapkit.foundation.traceflow.TraceFlow.dump`
+        * :term:`trace.tcp.packet`
+
+    """
+
+    #: Data link type from global header.
+    protocol: 'Enum_LinkType'
+    #: Frame number.
+    index: 'int'
+    #: Extracted frame info.
+    frame: 'Data_Frame | dict[str, Any]'
+    #: TCP synchronise (SYN) flag.
+    syn: 'bool'
+    #: TCP finish (FIN) flag.
+    fin: 'bool'
+    #: Source IP.
+    src: '_AT'
+    #: Destination IP.
+    dst: '_AT'
+    #: TCP source port.
+    srcport: 'int'
+    #: TCP destination port.
+    dstport: 'int'
+    #: Frame timestamp.
+    timestamp: 'float'
+
+    if TYPE_CHECKING:
+        def __init__(self, protocol: 'Enum_LinkType', index: 'int', frame: 'Data_Frame | dict[str, Any]', syn: 'bool', fin: 'bool', src: '_AT', dst: '_AT',
+                     srcport: 'int', dstport: 'int', timestamp: 'float') -> 'None': ...  # pylint: disable=unused-argument,super-init-not-called,multiple-statements,line-too-long
+
+
+@info_final
+class Buffer(Info):
+    """Data structure for **TCP flow tracing**.
+
+    See Also:
+        * :attr:`pcapkit.foundation.traceflow.TraceFlow.index`
+        * :term:`trace.tcp.buffer`
+
+    """
+
+    #: Output dumper object.
+    fpout: 'Dumper'
+    #: List of frame index.
+    index: 'list[int]'
+    #: Flow label generated from ``BUFID``.
+    label: 'str'
+
+    if TYPE_CHECKING:
+        def __init__(self, fpout: 'Dumper',
+                     index: 'list[int]', label: 'str') -> 'None': ...  # pylint: disable=unused-argument,super-init-not-called,multiple-statements
+
+
+@info_final
+class Index(Info):
+    """Data structure for **TCP flow tracing**.
+
+    See Also:
+        * element from :attr:`pcapkit.foundation.traceflow.TraceFlow.index`
+          *tuple*
+        * :term:`trace.tcp.index`
+
+    """
+
+    #: Output filename if exists.
+    fpout: 'Optional[str]'
+    #: Tuple of frame index.
+    index: 'tuple[int, ...]'
+    #: Flow label generated from ``BUFID``.
+    label: 'str'
+
+    if TYPE_CHECKING:
+        def __init__(self, fpout: 'Optional[str]', index: 'tuple[int, ...]',
+                     label: 'str') -> 'None': ...  # pylint: disable=unused-argument,super-init-not-called,multiple-statements

pcapkit/foundation/traceflow/tcp.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=import-outside-toplevel
+"""Follow TCP Streams
+========================
+
+.. module:: pcapkit.foundation.traceflow.tcp
+
+:mod:`pcapkit.foundation.traceflow.tcp` is the interface to trace
+TCP flows from a series of packets and connections.
+
+"""
+from typing import TYPE_CHECKING, Generic, overload
+
+from pcapkit.foundation.traceflow.data.tcp import _AT, Buffer, BufferID, Index, Packet
+from pcapkit.foundation.traceflow.traceflow import TraceFlowBase as TraceFlow
+from pcapkit.protocols.transport.tcp import TCP as TCP_Protocol
+
+__all__ = ['TCP']
+
+if TYPE_CHECKING:
+    from dictdumper.dumper import Dumper
+    from typing_extensions import Literal
+
+
+class TCP(TraceFlow[BufferID, Buffer, Index, Packet[_AT]], Generic[_AT]):
+    """Trace TCP flows.
+
+    Args:
+        fout: output path
+        format: output format
+        byteorder: output file byte order
+        nanosecond: output nanosecond-resolution file flag
+        *args: Arbitrary positional arguments.
+        **kwargs: Arbitrary keyword arguments.
+
+    """
+
+    ##########################################################################
+    # Defaults.
+    ##########################################################################
+
+    #: Protocol name of current reassembly object.
+    __protocol_name__ = 'TCP'
+    #: Protocol of current reassembly object.
+    __protocol_type__ = TCP_Protocol
+
+    ##########################################################################
+    # Methods.
+    ##########################################################################
+
+    def dump(self, packet: 'Packet[_AT]') -> 'None':
+        """Dump frame to output files.
+
+        Arguments:
+            packet: a flow packet (:term:`trace.tcp.packet`)
+
+        """
+        # fetch flow label
+        output = self.trace(packet, output=True)
+
+        # dump files
+        output(packet.frame, name=f'Frame {packet.index}')  # pylint: disable=not-callable
+
+    @overload
+    def trace(self, packet: 'Packet[_AT]', *, output: 'Literal[True]' = ...) -> 'Dumper': ...
+    @overload
+    def trace(self, packet: 'Packet[_AT]', *, output: 'Literal[False]' = ...) -> 'str': ...
+
+    def trace(self, packet: 'Packet[_AT]', *, 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.
+
+        Notes:
+            The flow label is formatted as following:
+
+            .. code-block:: python
+
+               f'{packet.src}_{packet.srcport}-{packet.dst}_{info.dstport}-{packet.timestamp}'
+
+        """
+        # clear cache
+        self.__cached__['submit'] = None
+
+        # Buffer Identifier
+        BUFID = (packet.src, packet.srcport, packet.dst, packet.dstport)  # type: BufferID
+        # SYN = packet.syn  # Synchronise Flag (Establishment)
+        FIN = packet.fin  # Finish Flag (Termination)
+
+        # # when SYN is set, reset buffer of this seesion
+        # if SYN and BUFID in self._buffer:
+        #     temp = self._buffer.pop(BUFID)
+        #     temp['fpout'] = (self._fproot, self._fdpext)
+        #     temp['index'] = tuple(temp['index'])
+        #     self._stream.append(Info(temp))
+
+        # initialise buffer with BUFID
+        if BUFID not in self._buffer:
+            if packet.src.version == 4:
+                label = f'{packet.src}_{packet.srcport}-{packet.dst}_{packet.dstport}-{packet.timestamp}'
+            else:
+                label = f'{packet.src}_{packet.srcport}-{packet.dst}_{packet.dstport}-{packet.timestamp}'.replace(':', '.')
+            self._buffer[BUFID] = Buffer(
+                fpout=self._foutio(fname=f'{self._fproot}/{label}{self._fdpext or ""}', protocol=packet.protocol,
+                                   byteorder=self._endian, nanosecond=self._nnsecd),
+                index=[],
+                label=label,
+            )
+
+        # trace frame record
+        self._buffer[BUFID].index.append(packet.index)
+        fpout = self._buffer[BUFID].fpout
+        label = self._buffer[BUFID].label
+
+        # when FIN is set, submit buffer of this session
+        if FIN:
+            buf = self._buffer.pop(BUFID)
+            # fpout, label = buf['fpout'], buf['label']
+
+            index = Index(
+                fpout=f'{self._fproot}/{label}{self._fdpext}' if self._fdpext is not None else None,
+                index=tuple(buf.index),
+                label=label,
+            )
+            for callback in self.__callback_fn__:
+                callback(index)
+            self._stream.append(index)
+
+        # return label or output object
+        return fpout if output else label
+
+    def submit(self) -> 'tuple[Index, ...]':
+        """Submit traced TCP flows.
+
+        Returns:
+            Traced TCP flow (:term:`trace.tcp.index`).
+
+        """
+        if (cached := self.__cached__.get('submit')) is not None:
+            return cached
+
+        ret = []  # type: list[Index]
+        for buf in self._buffer.values():
+            ret.append(Index(fpout=f"{self._fproot}/{buf.label}{self._fdpext}" if self._fdpext else None,
+                             index=tuple(buf.index),
+                             label=buf.label,))
+        ret.extend(self._stream)
+        ret_submit = tuple(ret)
+
+        self.__cached__['submit'] = ret_submit
+        return ret_submit