pypcapkit 1.3.3.post1__cp313-none-any.whl

pcapkit/foundation/reassembly/reassembly.py

@@ -0,0 +1,389 @@
+# -*- coding: utf-8 -*-
+"""Base Class
+================
+
+.. module:: pcapkit.foundation.reassembly.reassembly
+
+:mod:`pcapkit.foundation.reassembly.reassembly` contains
+:class:`~pcapkit.foundation.reassembly.reassembly.Reassembly` only,
+which is an abstract base class for all reassembly classes,
+bases on algorithms described in :rfc:`791` and :rfc:`815`,
+implements datagram reassembly of IP and TCP packets.
+
+"""
+import abc
+from typing import TYPE_CHECKING, Generic, Type, TypeVar, cast
+
+from pcapkit.protocols import __proto__ as protocol_registry
+from pcapkit.protocols.misc.raw import Raw
+from pcapkit.utilities.exceptions import UnsupportedCall
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional, Type
+
+    from typing_extensions import Self
+
+    from pcapkit.corekit.infoclass import Info
+    from pcapkit.protocols.protocol import ProtocolBase as Protocol
+
+    CallbackFn = Callable[[list['_DT']], None]
+
+__all__ = ['Reassembly']
+
+# packet
+_PT = TypeVar('_PT', bound='Info')
+# datagram
+_DT = TypeVar('_DT', bound='Info')
+# buffer ID
+_IT = TypeVar('_IT', bound='tuple')
+# buffer
+_BT = TypeVar('_BT', bound='Info')
+
+
+class ReassemblyMeta(abc.ABCMeta):
+    """Meta class to add dynamic support to :class:`Reassembly`.
+
+    This meta class is used to generate necessary attributes for the
+    :class:`Reassembly` class. It can be useful to reduce unnecessary
+    registry calls and simplify the customisation process.
+
+    """
+    if TYPE_CHECKING:
+        #: Protocol name of current reassembly object.
+        __protocol_name__: 'str'
+        #: Protocol of current reassembly object.
+        __protocol_type__: 'Type[Protocol]'
+
+    @property
+    def name(cls) -> 'str':
+        """Protocol name of current reassembly object."""
+        if hasattr(cls, '__protocol_name__'):
+            return cls.__protocol_name__
+        return cls.__name__
+
+    @property
+    def protocol(cls) -> 'Type[Protocol]':
+        """Protocol of current reassembly object."""
+        if hasattr(cls, '__protocol_type__'):
+            return cls.__protocol_type__
+        return protocol_registry.get(cls.name.upper(), Raw)
+
+
+class ReassemblyBase(Generic[_PT, _DT, _IT, _BT], metaclass=ReassemblyMeta):
+    """Base class for reassembly procedure.
+
+    Args:
+        strict: if return all datagrams (including those not
+                implemented) when submit
+        store: if store reassembled datagram in memory, i.e.,
+            :attr:`self._dtgram <_dtgram>` (if not, datagram
+            will be discarded after callback)
+
+    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]'
+
+    _flag_s: 'bool'
+    _flag_d: 'bool'
+    _flag_n: 'bool'
+
+    # Internal data storage for cached properties.
+    __cached__: 'dict[str, Any]'
+
+    ##########################################################################
+    # Properties.
+    ##########################################################################
+
+    @property
+    def name(self) -> 'str':
+        """Protocol name of current reassembly 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 reassembly 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]
+
+    # total number of reassembled packets
+    @property
+    def count(self) -> 'int':
+        """Total number of reassembled packets."""
+        if self._flag_n:
+            self.__cached__.clear()
+            self._flag_n = False
+
+        if (cached := self.__cached__.get('count')) is not None:
+            return cached
+
+        ret = len(self.datagram)
+        self.__cached__['count'] = ret
+        return ret
+
+    # reassembled datagram
+    @property
+    def datagram(self) -> 'tuple[_DT, ...]':
+        """Reassembled datagram.
+
+        Raises:
+            UnsupportedCall: If :attr:`self._flag_d <_flag_d>` is
+                set to :data:`False`.
+
+        """
+        if not self._flag_d:
+            raise UnsupportedCall(f'{self.__class__.__name__}(store=False) has no attribute "datagram"')
+
+        if self._buffer:
+            return self.fetch()
+        return tuple(self._dtgram)
+
+    ##########################################################################
+    # Methods.
+    ##########################################################################
+
+    # reassembly procedure
+    @abc.abstractmethod
+    def reassembly(self, info: '_PT') -> 'None':
+        """Reassembly procedure.
+
+        Arguments:
+            info: info dict of packets to be reassembled
+
+        """
+        # clear cache
+        self._flag_n = False
+        self.__cached__['count'] = None
+        self.__cached__['fetch'] = None
+
+    # submit reassembled payload
+    @abc.abstractmethod
+    def submit(self, buf: '_BT', **kwargs: 'Any') -> 'list[_DT]':
+        """Submit reassembled payload.
+
+        Arguments:
+            buf: buffer dict of reassembled packets
+            **kwargs: arbitrary keyword arguments
+
+        """
+
+    # fetch datagram
+    def fetch(self) -> 'tuple[_DT, ...]':
+        """Fetch datagram.
+
+        Returns:
+            Tuple of reassembled datagrams.
+
+        Fetch reassembled datagrams from
+        :attr:`self._dtgram <pcapkit.foundation.reassembly.reassembly.Reassembly._dtgram>`
+        and returns a *tuple* of such datagrams.
+
+        If no cache found, the method will call
+        :meth:`self.submit <pcapkit.foundation.reassembly.reassembly.Reassembly.submit>`
+        to *forcedly* obtain newly reassembled payload. Otherwise, the
+        already calculated
+        :attr:`self._dtgram <pcapkit.foundation.reassembly.reassembly.Reassembly._dtgram>`
+        will be returned.
+
+        """
+        if self._flag_n:
+            self.__cached__.clear()
+            self._flag_n = False
+
+        if (cached := self.__cached__.get('fetch')) is not None:
+            return cached
+
+        temp_dtgram = []  # type: list[_DT]
+        for (bufid, buffer) in self._buffer.items():
+            temp_dtgram.extend(
+                self.submit(buffer, bufid=bufid)
+            )
+        temp_dtgram.extend(self._dtgram)
+        ret = tuple(temp_dtgram)
+
+        self.__cached__['fetch'] = ret
+        return ret
+
+    # return datagram index
+    def index(self, pkt_num: 'int') -> 'Optional[int]':
+        """Return datagram index.
+
+        Arguments:
+            pkt_num: index of packet
+
+        Returns:
+            Reassembled datagram index which was from No. ``pkt_num`` packet;
+            if not found, returns :obj:`None`.
+
+        """
+        for counter, datagram in enumerate(self.datagram):
+            if pkt_num in datagram.index:  # type: ignore[attr-defined]
+                return counter
+        return None
+
+    # run automatically
+    def run(self, packets: 'list[_PT]') -> 'None':
+        """Run automatically.
+
+        Arguments:
+            packets: list of packet dicts to be reassembled
+
+        """
+        for packet in packets:
+            self.reassembly(packet)
+
+    # register callback function
+    @classmethod
+    def register(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 to be inserted in the callback list,; by
+                default, the callback will be appended to the end
+                of the list
+
+        """
+        if index is not None:
+            cls.__callback_fn__.insert(index, callback)
+        else:
+            cls.__callback_fn__.append(callback)
+
+    ##########################################################################
+    # 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, *, strict: 'bool' = True, store: 'bool' = True) -> 'None':
+        """Initialise packet reassembly.
+
+        Args:
+            strict: if return all datagrams (including those not
+                implemented) when submit
+            store: if store reassembled datagram in memory, i.e.,
+                :attr:`self._dtgram <_dtgram>` (if not, datagram
+                will be discarded after callback)
+
+        """
+        #: bool: Strict mode flag. If set to :data:`True`, all
+        #: data will be returned, including those not completely
+        #: reassembled; otherwise, only completely reassembled
+        #: data will be returned.
+        self._flag_s = strict
+        #: bool: Store mode flag. If set to :data:`True`, all
+        #: reassembled datagram will be stored in memory, i.e.,
+        #: :attr:`self._dtgram <_dtgram>`; otherwise, datagram
+        #: will be discarded after callback.
+        self._flag_d = store
+        #: bool: New datagram flag. If set to :data:`True`, the
+        #: :attr:`self._dtgram <_dtgram>` will be repopulated.
+        self._flag_n = False
+
+        #: dict[_IT, _BT]: Dict buffer field. This field is used to
+        #: store reassembled packets in the form of ``{bufid: buffer}``.
+        self._buffer = {}  # type: dict[_IT, _BT]
+        #: list[_DT]: List reassembled datagram. This list is used
+        #: to store reassembled datagrams.
+        self._dtgram = []  # type: list[_DT]
+
+    def __call__(self, packet: '_PT') -> 'None':
+        """Call packet reassembly.
+
+        Arguments:
+            packet: packet dict to be reassembled
+                (detailed format described in corresponding protocol)
+
+        """
+        self._flag_n = True
+        self.reassembly(packet)
+
+    def __init_subclass__(cls) -> 'None':
+        """Initialise subclass.
+
+        This method is to be used for generating necessary attributes
+        for the :class:`Reassembly` class. It can be useful to reduce
+        unnecessary registry calls and simplify the customisation process.
+
+        """
+        cls.__callback_fn__ = []
+
+
+class Reassembly(ReassemblyBase[_PT, _DT, _IT, _BT], Generic[_PT, _DT, _IT, _BT]):
+    """Base flow tracing class.
+
+    Example:
+
+        Use keyword argument ``protocol`` to specify the protocol
+        name at class definition:
+
+        .. code-block:: python
+
+           class MyProtocol(Reassembly, protocol='my_protocol'):
+               ...
+
+    Arguments:
+        strict: if return all datagrams (including those not
+                implemented) when submit
+        store: if store reassembled datagram in memory, i.e.,
+            :attr:`self._dtgram <_dtgram>` (if not, datagram
+            will be discarded after callback)
+
+    """
+
+    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_reassembly`.
+
+        """
+        if protocol is None:
+            protocol = cast('str', cls.name)
+
+        from pcapkit.foundation.extraction import Extractor
+        Extractor.register_reassembly(protocol.lower(), cls)
+
+        return super().__init_subclass__()

pcapkit/foundation/reassembly/tcp.py

@@ -0,0 +1,249 @@
+# -*- coding: utf-8 -*-
+"""TCP Datagram Reassembly
+=============================
+
+.. module:: pcapkit.foundation.reassembly.tcp
+
+:mod:`pcapkit.foundation.reassembly.tcp` contains
+:class:`~pcapkit.foundation.reassembly.reassembly.Reassembly` only,
+which reconstructs fragmented TCP packets back to origin.
+
+"""
+import sys
+from typing import TYPE_CHECKING
+
+from pcapkit.foundation.reassembly.data.tcp import (Buffer, BufferID, Datagram, DatagramID,
+                                                    Fragment, HoleDiscriptor, Packet)
+from pcapkit.foundation.reassembly.reassembly import ReassemblyBase as Reassembly
+from pcapkit.protocols.transport.tcp import TCP as TCP_Protocol
+
+if TYPE_CHECKING:
+    from typing import Type
+
+__all__ = ['TCP']
+
+
+class TCP(Reassembly[Packet, Datagram, BufferID, Buffer]):
+    """Reassembly for TCP payload.
+
+    Args:
+        strict: if return all datagrams (including those not
+                implemented) when submit
+        store: if store reassembled datagram in memory, i.e.,
+            :attr:`self._dtgram <pcapkit.foundation.reassembly.reassembly.Reassembly._dtgram>`
+            (if not, datagram will be discarded after callback)
+
+    Example:
+        >>> from pcapkit.foundation.reassembly import TCP
+        # Initialise instance:
+        >>> tcp_reassembly = TCP()
+        # Call reassembly:
+        >>> tcp_reassembly(packet_dict)
+        # Fetch result:
+        >>> result = tcp_reassembly.datagram
+
+    """
+    if TYPE_CHECKING:
+        protocol: 'Type[TCP_Protocol]'
+
+    ##########################################################################
+    # Defaults.
+    ##########################################################################
+
+    #: Protocol name of current reassembly object.
+    __protocol_name__ = 'TCP'
+    #: Protocol of current reassembly object.
+    __protocol_type__ = TCP_Protocol
+
+    ##########################################################################
+    # Methods.
+    ##########################################################################
+
+    def reassembly(self, info: 'Packet') -> 'None':
+        """Reassembly procedure.
+
+        Arguments:
+            info: :term:`info <reasm.tcp.packet>` dict of packets to be reassembled
+
+        """
+        # clear cache
+        self._flag_n = False
+        self.__cached__.clear()
+
+        BUFID = info.bufid  # Buffer Identifier
+        DSN = info.dsn      # Data Sequence Number
+        ACK = info.ack      # Acknowledgement Number
+        FIN = info.fin      # Finish Flag (Termination)
+        RST = info.rst      # Reset Connection Flag (Termination)
+        SYN = info.syn      # Synchronise Flag (Establishment)
+
+        # when SYN is set, reset buffer of existing session
+        if SYN and BUFID in self._buffer:
+            self._dtgram.extend(
+                self.submit(self._buffer.pop(BUFID), bufid=BUFID)
+            )
+
+        # initialise buffer with BUFID & ACK
+        if BUFID not in self._buffer:
+            self._buffer[BUFID] = Buffer(
+                hdl=[
+                    HoleDiscriptor(
+                        first=info.len,
+                        last=sys.maxsize,
+                    ),
+                ],
+                hdr=info.header if SYN else b'',
+                ack={
+                    ACK: Fragment(
+                        ind=[
+                            info.num,
+                        ],
+                        isn=info.dsn,
+                        len=info.len,
+                        raw=info.payload,
+                    ),
+                },
+            )
+        else:
+            # initialise buffer with ACK
+            if ACK not in self._buffer[BUFID].ack:
+                self._buffer[BUFID].ack[ACK] = Fragment(
+                    ind=[
+                        info.num,
+                    ],
+                    isn=info.dsn,
+                    len=info.len,
+                    raw=info.payload,
+                )
+            else:
+                # put header into header buffer
+                if SYN:
+                    self._buffer[BUFID].__update__(hdr=info.header)
+
+                # append packet index
+                self._buffer[BUFID].ack[ACK].ind.append(info.num)
+
+                # record fragment payload
+                ISN = self._buffer[BUFID].ack[ACK].isn       # Initial Sequence Number
+                RAW = self._buffer[BUFID].ack[ACK].raw       # Raw Payload Data
+                if DSN >= ISN:  # if fragment goes after existing payload
+                    LEN = self._buffer[BUFID].ack[ACK].len
+                    GAP = DSN - (ISN + LEN)     # gap length between payloads
+                    if GAP >= 0:    # if fragment goes after existing payload
+                        RAW += bytearray(GAP) + info.payload
+                    else:           # if fragment partially overlaps existing payload
+                        RAW[DSN-ISN:] = info.payload
+                else:           # if fragment exceeds existing payload
+                    LEN = info.len
+                    GAP = ISN - (DSN + LEN)     # gap length between payloads
+                    self._buffer[BUFID].ack[ACK].__update__(
+                        isn=DSN,
+                    )
+                    if GAP >= 0:    # if fragment exceeds existing payload
+                        RAW = info.payload + bytearray(GAP) + RAW
+                    else:           # if fragment partially overlaps existing payload
+                        RAW = info.payload + RAW[ISN-GAP:]
+                #self._buffer[BUFID].ack[ACK].raw = RAW       # update payload datagram
+                #self._buffer[BUFID].ack[ACK].len = len(RAW)  # update payload length
+                self._buffer[BUFID].ack[ACK].__update__(
+                    raw=RAW,       # update payload datagram
+                    len=len(RAW),  # update payload length
+                )
+
+            # update hole descriptor list
+            HDL = self._buffer[BUFID].hdl                          # HDL alias
+            for (index, hole) in enumerate(HDL):                   # step one
+                if info.first > hole.last:                         # step two
+                    continue
+                if info.last < hole.first:                         # step three
+                    continue
+                del HDL[index]                                     # step four
+                if info.first > hole.first:                        # step five
+                    new_hole = HoleDiscriptor(
+                        first=hole.first,
+                        last=info.first - 1,
+                    )
+                    HDL.insert(index, new_hole)
+                    index += 1
+                if info.last < hole.last and not FIN and not RST:  # step six
+                    new_hole = HoleDiscriptor(
+                        first=info.last + 1,
+                        last=hole.last
+                    )
+                    HDL.insert(index, new_hole)
+                break                                              # step seven
+            #self._buffer[BUFID].hdl = HDL                         # update HDL
+
+        # when FIN/RST is set, submit buffer of this session
+        if FIN or RST:
+            self._dtgram.extend(
+                self.submit(self._buffer.pop(BUFID), bufid=BUFID)
+            )
+
+    def submit(self, buf: 'Buffer', *, bufid: 'BufferID') -> 'list[Datagram]':  # type: ignore[override] # pylint: disable=arguments-differ
+        """Submit reassembled payload.
+
+        Arguments:
+            buf: :term:`buffer <reasm.tcp.buffer>` dict of reassembled packets
+            bufid: buffer identifier
+
+        Returns:
+            Reassembled :term:`packets <reasm.tcp.datagram>`.
+
+        """
+        datagram = []  # type: list[Datagram] # reassembled datagram
+        HDL = buf.hdl                         # hole descriptor list
+
+        # check through every buffer with ACK
+        for (ack, buffer) in buf.ack.items():
+            # if this buffer is not implemented
+            # go through every hole and extract received payload
+            if len(HDL) > 2 and self._flag_s:
+                data = []  # type: list[bytes]
+                start = stop = 0
+                for hole in HDL:
+                    stop = hole.first
+                    byte = buffer.raw[start:stop]
+                    start = hole.last
+                    if byte:    # strip empty payload
+                        data.append(byte)
+                byte = buffer.raw[start:]
+                if byte:    # strip empty payload
+                    data.append(bytes(byte))
+                if data:    # strip empty buffer
+                    packet = Datagram(
+                        completed=False,
+                        id=DatagramID(
+                            src=(bufid[0], bufid[1]),
+                            dst=(bufid[2], bufid[3]),
+                            ack=ack,
+                        ),
+                        index=tuple(buffer.ind),
+                        header=buf.hdr,
+                        payload=tuple(data),
+                        packet=None,
+                    )
+                    datagram.append(packet)
+
+            # if this buffer is implemented
+            # export payload data & convert into bytes
+            else:
+                payload = buffer.raw
+                if payload:    # strip empty buffer
+                    packet = Datagram(
+                        completed=True,
+                        id=DatagramID(
+                            src=(bufid[0], bufid[1]),
+                            dst=(bufid[2], bufid[3]),
+                            ack=ack,
+                        ),
+                        index=tuple(buffer.ind),
+                        header=buf.hdr,
+                        payload=bytes(payload),
+                        packet=self.protocol.analyze((bufid[1], bufid[3]), bytes(payload)),
+                    )
+                    datagram.append(packet)
+
+        for callback in self.__callback_fn__:
+            callback(datagram)
+        return datagram

pcapkit/foundation/registry/__init__.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+"""Registry Management
+=========================
+
+.. module:: pcapkit.foundation.registry
+
+This module provides the registry management for :mod:`pcapkit`, as the module
+contains various registry points.
+
+"""
+
+from pcapkit.foundation.registry.foundation import *
+from pcapkit.foundation.registry.protocols import *
+
+__all__ = [
+    'register_extractor_engine',
+
+    'register_dumper',
+    'register_extractor_dumper', 'register_traceflow_dumper',
+
+    'register_protocol',
+
+    'register_linktype',
+    'register_pcap', 'register_pcapng',
+
+    'register_ethertype',
+
+    'register_transtype',
+    'register_ipv4_option', 'register_hip_parameter', 'register_hopopt_option',
+    'register_ipv6_opts_option', 'register_ipv6_route_routing',
+    'register_mh_message', 'register_mh_option', 'register_mh_extension',
+
+    'register_apptype',
+    'register_tcp', 'register_udp',
+    'register_tcp_option', 'register_tcp_mp_option',
+
+    'register_http_frame',
+
+    'register_pcapng_block', 'register_pcapng_option', 'register_pcapng_secrets',
+    'register_pcapng_record',
+]