pypcapkit 1.3.3.post1__cp313-none-any.whl

pcapkit/protocols/transport/transport.py

@@ -0,0 +1,158 @@
+# -*- coding: utf-8 -*-
+"""Base Protocol
+===================
+
+.. module:: pcapkit.protocols.transport.transport
+
+:mod:`pcapkit.protocols.transport.transport` contains
+:class:`~pcapkit.protocols.transport.transport.Transport`,
+which is a base class for transport layer protocols, eg.
+:class:`~pcapkit.protocols.transport.transport.tcp.TCP` and
+:class:`~pcapkit.protocols.transport.transport.udp.UDP`.
+
+"""
+import io
+from typing import TYPE_CHECKING, Generic
+
+from pcapkit.corekit.module import ModuleDescriptor
+from pcapkit.protocols.protocol import _PT, _ST
+from pcapkit.protocols.protocol import ProtocolBase as Protocol
+from pcapkit.utilities.exceptions import StructError, UnsupportedCall, stacklevel
+from pcapkit.utilities.logging import DEVMODE, logger
+from pcapkit.utilities.warnings import RegistryWarning, warn
+
+if TYPE_CHECKING:
+    from typing import Any, DefaultDict, Optional, Type
+
+    from typing_extensions import Literal
+
+__all__ = ['Transport']
+
+
+class Transport(Protocol[_PT, _ST], Generic[_PT, _ST]):  # pylint: disable=abstract-method
+    """Abstract base class for transport layer protocol family."""
+
+    if TYPE_CHECKING:
+        #: Protocol index mapping for decoding next layer,
+        #: c.f. :meth:`self._decode_next_layer <pcapkit.protocols.transport.transport.Transport._decode_next_layer>`
+        #: & :meth:`self._import_next_layer <pcapkit.protocols.protocol.Protocol._import_next_layer>`.
+        __proto__: 'DefaultDict[int, ModuleDescriptor[Protocol] | Type[Protocol]]'
+
+    ##########################################################################
+    # Defaults.
+    ##########################################################################
+
+    #: Layer of protocol.
+    __layer__ = 'Transport'  # type: Literal['Transport']
+
+    ##########################################################################
+    # Properties.
+    ##########################################################################
+
+    # protocol layer
+    @property
+    def layer(self) -> 'Literal["Transport"]':
+        """Protocol layer."""
+        return self.__layer__
+
+    ##########################################################################
+    # Methods.
+    ##########################################################################
+
+    @classmethod
+    def register(cls, code: 'int', protocol: 'ModuleDescriptor[Protocol] | Type[Protocol]') -> 'None':
+        """Register a new protocol class.
+
+        Notes:
+            The full qualified class name of the new protocol class
+            should be as ``{protocol.module}.{protocol.name}``.
+
+        Arguments:
+            code: port number
+            protocol: module name
+
+        Important:
+            This method must be called from a non-abstract class, as the
+            protocol map should be associated directly with specific
+            transport layer protocol type.
+
+        """
+        if cls is Transport:
+            raise UnsupportedCall(f'{cls.__name__} is an abstract class')
+
+        if isinstance(protocol, ModuleDescriptor):
+            protocol = protocol.klass
+        if not issubclass(protocol, Protocol):
+            raise TypeError(f'protocol must be a Protocol subclass, not {protocol!r}')
+        if code in cls.__proto__:
+            warn(f'port {code} already registered, overwriting', RegistryWarning)
+        cls.__proto__[code] = protocol
+
+    @classmethod
+    def analyze(cls, ports: 'tuple[int, int]', payload: 'bytes', **kwargs: 'Any') -> 'Protocol':  # type: ignore[override] # pylint: disable=arguments-renamed
+        """Analyse packet payload.
+
+        Args:
+            ports: Source & destination port numbers.
+            payload: Packet payload.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Parsed payload as a :class:`~pcapkit.protocols.protocol.Protocol`
+            instance.
+
+        """
+        if ports[0] in cls.__proto__:
+            protocol = cls.__proto__[ports[0]]
+        else:
+            protocol = cls.__proto__[ports[1]]
+
+        if isinstance(protocol, ModuleDescriptor):
+            protocol = protocol.klass
+
+        payload_io = io.BytesIO(payload)
+        try:
+            report = protocol(payload_io, len(payload), **kwargs)  # type: ignore[abstract]
+        except Exception as exc:
+            if isinstance(exc, StructError) and exc.eof:  # pylint: disable=no-member
+                from pcapkit.protocols.misc.null import NoPayload as protocol  # pylint: disable=import-outside-toplevel # isort:skip
+            else:
+                from pcapkit.protocols.misc.raw import Raw as protocol  # pylint: disable=import-outside-toplevel # isort:skip
+            # error = traceback.format_exc(limit=1).strip().rsplit(os.linesep, maxsplit=1)[-1]
+
+            # log error
+            logger.error(str(exc), exc_info=exc, stack_info=DEVMODE, stacklevel=stacklevel())
+
+            report = protocol(payload_io, len(payload), **kwargs)  # type: ignore[abstract]
+        return report
+
+    ##########################################################################
+    # Utilities.
+    ##########################################################################
+
+    def _decode_next_layer(self, dict_: '_PT', ports: 'tuple[int, int]', length: 'Optional[int]' = None, *,  # type: ignore[override]
+                           packet: 'Optional[dict[str, Any]]' = None) -> '_PT':  # pylint: disable=arguments-renamed
+        """Decode next layer protocol.
+
+        The method will check if the next layer protocol is supported based on
+        the source and destination port numbers. We will use the lower port
+        number from both ports as the primary key to lookup the next layer.
+
+        Arguments:
+            dict_: info buffer
+            ports: source & destination port numbers
+            length: valid (*non-padding*) length
+            packet: packet info (passed from :meth:`self.unpack <pcapkit.protocols.protocol.Protocol.unpack>`)
+
+        Returns:
+            Current protocol with next layer extracted.
+
+        """
+        sort_port = sorted(ports)
+        if sort_port[0] in self.__proto__:
+            proto = sort_port[0]
+        elif sort_port[1] in self.__proto__:
+            proto = sort_port[1]
+        else:
+            proto = None
+        return super()._decode_next_layer(dict_, proto, length, packet=packet)  # type: ignore[arg-type]

pcapkit/protocols/transport/udp.py

@@ -0,0 +1,214 @@
+# -*- coding: utf-8 -*-
+# mypy: disable-error-code=dict-item
+"""UDP - User Datagram Protocol
+==================================
+
+.. module:: pcapkit.protocols.transport.udp
+
+:mod:`pcapkit.protocols.transport.udp` contains
+:class:`~pcapkit.protocols.transport.udp.UDP` only,
+which implements extractor for User Datagram Protocol
+(UDP) [*]_, whose structure is described as below:
+
+======= ========= ===================== ===============================
+Octets      Bits        Name                    Description
+======= ========= ===================== ===============================
+  0           0   ``udp.srcport``             Source Port
+  2          16   ``udp.dstport``             Destination Port
+  4          32   ``udp.len``                 Length (header includes)
+  6          48   ``udp.checksum``            Checksum
+======= ========= ===================== ===============================
+
+.. [*] https://en.wikipedia.org/wiki/User_Datagram_Protocol
+
+"""
+import collections
+from typing import TYPE_CHECKING
+
+from pcapkit.const.reg.transtype import TransType as Enum_TransType
+from pcapkit.corekit.module import ModuleDescriptor
+from pcapkit.protocols.data.transport.udp import UDP as Data_UDP
+from pcapkit.protocols.schema.transport.udp import UDP as Schema_UDP
+from pcapkit.protocols.transport.transport import Transport
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+    from typing_extensions import Literal
+
+    from pcapkit.const.reg.apptype import AppType as Enum_AppType
+    from pcapkit.protocols.protocol import ProtocolBase as Protocol
+    from pcapkit.protocols.schema.schema import Schema
+
+__all__ = ['UDP']
+
+
+class UDP(Transport[Data_UDP, Schema_UDP],
+          schema=Schema_UDP, data=Data_UDP):
+    """This class implements User Datagram Protocol.
+
+    This class currently supports parsing of the following protocols, which are
+    registered in the :attr:`self.__proto__ <pcapkit.protocols.transport.udp.UDP.__proto__>`
+    attribute:
+
+    .. list-table::
+       :header-rows: 1
+
+       * - Port Number
+         - Protocol
+       * - 80
+         - :class:`pcapkit.protocols.application.http.HTTP`
+
+    """
+
+    ##########################################################################
+    # Defaults.
+    ##########################################################################
+
+    #: DefaultDict[int, ModuleDescriptor[Protocol] | Type[Protocol]]: Protocol
+    #: index mapping for decoding next layer, c.f.
+    #: :meth:`self._decode_next_layer <pcapkit.protocols.transport.transport.Transport._decode_next_layer>`
+    #: & :meth:`self._import_next_layer <pcapkit.protocols.protocol.Protocol._import_next_layer>`.
+    __proto__ = collections.defaultdict(
+        lambda: ModuleDescriptor('pcapkit.protocols.misc.raw', 'Raw'),
+        {
+            80: ModuleDescriptor('pcapkit.protocols.application.http', 'HTTP'),  # HTTP
+        },
+    )
+
+    ##########################################################################
+    # Properties.
+    ##########################################################################
+
+    @property
+    def name(self) -> 'Literal["User Datagram Protocol"]':
+        """Name of current protocol."""
+        return 'User Datagram Protocol'
+
+    @property
+    def length(self) -> 'Literal[8]':
+        """Header length of current protocol."""
+        return 8
+
+    @property
+    def src(self) -> 'Enum_AppType':
+        """Source port."""
+        return self._info.srcport
+
+    @property
+    def dst(self) -> 'Enum_AppType':
+        """Destination port."""
+        return self._info.dstport
+
+    ##########################################################################
+    # Methods.
+    ##########################################################################
+
+    def read(self, length: 'Optional[int]' = None, **kwargs: 'Any') -> 'Data_UDP':  # pylint: disable=unused-argument
+        """Read User Datagram Protocol (UDP).
+
+        Structure of UDP header [:rfc:`768`]:
+
+        .. code-block:: text
+
+            0      7 8     15 16    23 24    31
+           +--------+--------+--------+--------+
+           |     Source      |   Destination   |
+           |      Port       |      Port       |
+           +--------+--------+--------+--------+
+           |                 |                 |
+           |     Length      |    Checksum     |
+           +--------+--------+--------+--------+
+           |
+           |          data octets ...
+           +---------------- ...
+
+        Args:
+            length: Length of packet data.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Parsed packet data.
+
+        """
+        if length is None:
+            length = len(self)
+        schema = self.__header__
+
+        udp = Data_UDP(
+            srcport=schema.srcport,
+            dstport=schema.dstport,
+            len=schema.len,
+            checksum=schema.checksum,
+        )
+
+        return self._decode_next_layer(udp, (udp.srcport.port, udp.dstport.port), udp.len - 8)
+
+    def make(self,
+             srcport: 'Enum_AppType | int' = 0,
+             dstport: 'Enum_AppType | int' = 0,
+             checksum: 'bytes' = b'\x00\x00',
+             payload: 'bytes | Schema | Protocol' = b'',
+             **kwargs: 'Any') -> 'Schema_UDP':
+        """Make (construct) packet data.
+
+        Args:
+            srcport: Source port.
+            dstport: Destination port.
+            checksum: Checksum.
+            payload: Payload data.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Constructed packet data.
+
+        """
+        return Schema_UDP(
+            srcport=srcport,
+            dstport=dstport,
+            len=8 + len(payload),
+            checksum=checksum,
+            payload=payload,
+        )
+
+    ##########################################################################
+    # Data models.
+    ##########################################################################
+
+    def __length_hint__(self) -> 'Literal[8]':
+        """Return an estimated length for the object."""
+        return 8
+
+    @classmethod
+    def __index__(cls) -> 'Enum_TransType':  # pylint: disable=invalid-index-returned
+        """Numeral registry index of the protocol.
+
+        Returns:
+            Numeral registry index of the protocol in `IANA`_.
+
+        .. _IANA: https://www.iana.org/assignments/protocol-numbers/protocol-numbers.xhtml
+
+        """
+        return Enum_TransType.UDP  # type: ignore[return-value]
+
+    ##########################################################################
+    # Utilities.
+    ##########################################################################
+
+    @classmethod
+    def _make_data(cls, data: 'Data_UDP') -> 'dict[str, Any]':  # type: ignore[override]
+        """Create key-value pairs from ``data`` for protocol construction.
+
+        Args:
+            data: protocol data
+
+        Returns:
+            Key-value pairs for protocol construction.
+
+        """
+        return {
+            'srcport': data.srcport,
+            'dstport': data.dstport,
+            'checksum': data.checksum,
+            'payload': cls._make_payload(data),
+        }

pcapkit/toolkit/__init__.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=unused-import, unused-wildcard-import
+"""Compatibility Tools
+=========================
+
+.. module:: pcapkit.toolkit
+
+:mod:`pcapkit.toolkit` provides several utility functions for
+compatibility of multiple engine support.
+
+"""
+# tools for default engine
+from pcapkit.toolkit.pcap import ipv4_reassembly, ipv6_reassembly, tcp_reassembly, tcp_traceflow
+from pcapkit.toolkit.pcapng import ipv4_reassembly as pcapng_ipv4_reassembly
+from pcapkit.toolkit.pcapng import ipv6_reassembly as pcapng_ipv6_reassembly
+from pcapkit.toolkit.pcapng import tcp_reassembly as pcapng_tcp_reassembly
+from pcapkit.toolkit.pcapng import tcp_traceflow as pcapng_tcp_traceflow
+from pcapkit.toolkit.pcapng import block2frame as pcapng_block2frame
+
+# # tools for DPKT engine
+# from pcapkit.toolkit.dpkt import ipv6_hdr_len as dpkt_ipv6_hdr_len
+# from pcapkit.toolkit.dpkt import packet2chain as dpkt_packet2chain
+# from pcapkit.toolkit.dpkt import packet2dict as dpkt_packet2dict
+# from pcapkit.toolkit.dpkt import ipv4_reassembly as dpkt_ipv4_reassembly
+# from pcapkit.toolkit.dpkt import ipv6_reassembly as dpkt_ipv6_reassembly
+# from pcapkit.toolkit.dpkt import tcp_reassembly as dpkt_tcp_reassembly
+# from pcapkit.toolkit.dpkt import tcp_traceflow as dpkt_tcp_traceflow
+
+# # tools for PyShark engine
+# from pcapkit.toolkit.pyshark import packet2dict as pyshark_packet2dict
+# from pcapkit.toolkit.pyshark import tcp_traceflow as pyshark_tcp_traceflow
+
+# # tools for Scapy engine
+# from pcapkit.toolkit.scapy import packet2chain as scapy_packet2chain
+# from pcapkit.toolkit.scapy import packet2dict as scapy_packet2dict
+# from pcapkit.toolkit.scapy import ipv4_reassembly as scapy_ipv4_reassembly
+# from pcapkit.toolkit.scapy import ipv6_reassembly as scapy_ipv6_reassembly
+# from pcapkit.toolkit.scapy import tcp_reassembly as scapy_tcp_reassembly
+# from pcapkit.toolkit.scapy import tcp_traceflow as scapy_tcp_traceflow
+
+__all__ = [
+    # default engine
+    'ipv4_reassembly', 'ipv6_reassembly', 'tcp_reassembly', 'tcp_traceflow',
+    'pcapng_ipv4_reassembly', 'pcapng_ipv6_reassembly', 'pcapng_tcp_reassembly', 'pcapng_tcp_traceflow',
+    'pcapng_block2frame',
+
+    # # DPKT engine
+    # 'dpkt_ipv6_hdr_len', 'dpkt_packet2chain', 'dpkt_packet2dict',
+    # 'dpkt_ipv4_reassembly', 'dpkt_ipv6_reassembly', 'dpkt_tcp_reassembly', 'dpkt_tcp_traceflow',
+
+    # # PyShark engine
+    # 'pyshark_packet2dict', 'pyshark_tcp_traceflow',
+
+    # # Scapy engine
+    # 'scapy_packet2chain', 'scapy_packet2dict',
+    # 'scapy_ipv4_reassembly', 'scapy_ipv6_reassembly', 'scapy_tcp_reassembly', 'scapy_tcp_traceflow',
+]