# -*- 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
import types
from typing import TYPE_CHECKING, Generic, TypeVar, 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, UnsupportedCall, stacklevel
from pcapkit.utilities.logging import get_logger
from pcapkit.utilities.warnings import FileWarning, FormatWarning, RegistryWarning, warn
__all__ = ['TraceFlow']
#: logging.Logger: Module-level logger, a child of the package-wide
#: :data:`pcapkit.utilities.logging.logger`.
logger = get_logger(__name__)
# NB: declared above the ``TYPE_CHECKING`` block, not below it, so that
# ``CallbackFn`` can name ``_IT`` outright. As a quoted forward reference it was
# resolvable only from this module's namespace, and every module that spells
# ``CallbackFn`` in an annotation -- ``pcapkit.foundation.registry.foundation``
# does -- has to evaluate the alias in its own.
_DT = TypeVar('_DT')
_BT = TypeVar('_BT', bound='Info')
_IT = TypeVar('_IT', bound='Info')
_PT = TypeVar('_PT', bound='Info')
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
CallbackFn = Callable[[_IT], None]
[docs]
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
bidirectional: trace both halves of a conversation as one flow
analyse: reassemble each flow's application layer
Note:
This class is for internal use only. For customisation, please use
:class:`TraceFlow` instead.
"""
if TYPE_CHECKING:
#: Protocol name of current flow tracing object.
__protocol_name__: 'str'
#: Protocol of current flow tracing object.
__protocol_type__: 'Type[ProtocolBase]'
#: List of callback functions upon a finalised flow.
__callback_fn__: 'list[CallbackFn]'
# Internal data storage for cached properties.
__cached__: 'dict[str, Any]'
##########################################################################
# Defaults.
##########################################################################
#: Format dumper mapping for writing output files. Each value pairs the dumper,
#: as a module descriptor or a :class:`dictdumper.dumper.Dumper` subclass, with
#: the file extension of its format.
__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', 'Tree'), '.txt'),
'txt': (ModuleDescriptor('dictdumper', 'Tree'), '.txt'),
},
) # type: DefaultDict[str, tuple[ModuleDescriptor[Dumper] | Type[Dumper], str | None]]
# Read-only copy of ``__output__`` as shipped. :meth:`register_dumper`
# handed the dumper a format shipped with restores this entry, so an
# override can be undone (:issue:`1363`).
_shipped_output = types.MappingProxyType(dict(__output__))
##########################################################################
# 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[ProtocolBase]':
"""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, ...]':
"""Index table for traced flows."""
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 fully qualified class name of the new dumper class
should be ``{dumper.module}.{dumper.name}``.
The overwrite guard fires only when the incumbent dumper differs
from the replacement, so re-registering the exact same object is
a silent no-op rather than a warning about nothing displaced --
the identity guard of the code-keyed registrars (:issue:`718`,
:issue:`739`). ``__output__``
maps each format to a ``(dumper, ext)`` pair, so the identity
check compares the incumbent *dumper* (index ``0``), not the
pair -- a re-registration that only changes ``ext`` is still
identity-equal on the dumper and stays silent, since the dumper
is what "the same object" means here, not the pair as a whole.
``__output__`` is also a :class:`collections.defaultdict`,
unlike the :class:`~pcapkit.foundation.extraction.Extractor` engine,
reassembly and flow tracing registries; :meth:`dict.get` does not invoke
the default factory the way ``cls.__output__[format]`` would, so
it stays non-inserting here as well.
The dumper ``format`` shipped with, as its descriptor or as the
class it names, restores the shipped entry rather than storing
the class, so an override can be undone and the deferred import
is kept (:issue:`1363`, :issue:`1364`).
Arguments:
format: format name
dumper: module descriptor or a :class:`dictdumper.dumper.Dumper` subclass
ext: file extension
Raises:
RegistryError: If ``dumper`` is not a class, or not a ``Dumper`` subclass.
Warns:
RegistryWarning: If a different dumper is already registered for
``format``; it is overwritten.
.. versionremoved:: 1.2.0
``register``, which this replaces with no alias.
.. versionchanged:: 1.3.0
``dumper``, a module descriptor or a class, replaces the ``module``
and ``class_`` name pair, with no alias.
"""
from pcapkit.interface.registry import ( # pylint: disable=import-outside-toplevel
restore_shipped_dumper, same_entry)
shipped = cls._shipped_output.get(format)
if isinstance(dumper, ModuleDescriptor) and (shipped is None or dumper != shipped[0]):
dumper = dumper.klass
if shipped is not None and same_entry(shipped[0], dumper):
restore_shipped_dumper(cls.__output__, format, shipped, ext)
return
if not isinstance(dumper, type):
raise RegistryError(f'dumper must be a class, not {dumper!r}')
if not issubclass(dumper, Dumper):
raise RegistryError(f'dumper must be a Dumper subclass, not {dumper!r}')
incumbent_entry = cls.__output__.get(format)
incumbent = incumbent_entry[0] if incumbent_entry is not None else None
if incumbent is not None and incumbent is not dumper:
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 a traced flow is finalised, with the flow's
index entry (:term:`trace.tcp.index`) 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)
@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 but is not a directory, and ``fmt``
is not supported.
Raises:
FileExists: If ``fout`` exists but is not a directory, and ``fmt``
is supported.
"""
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__)
dumper = make_dumper(output)
# NOTE: as above -- make_dumper()'s subclass is always called 'DictDumper'.
logger.debug('flow tracing output root %s, format %s via %s', fout, fmt, output.__name__)
return dumper, 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`).
"""
def finish(self) -> 'None':
"""Finalise every flow still being traced.
Called by :meth:`Extractor._cleanup
<pcapkit.foundation.extraction.Extractor._cleanup>` once the capture has
been read to its end, which is the point at which a flow that was never
superseded can be said to be over.
The base implementation does nothing, so a tracer that has no such notion
need not override it. :meth:`submit` must remain able to report a flow that
was never finalised, since nothing guarantees this is called: a tracer
driven directly rather than through an
:class:`~pcapkit.foundation.extraction.Extractor` never sees an end of
capture.
Implementations must be **idempotent**: :meth:`Extractor._cleanup
<pcapkit.foundation.extraction.Extractor._cleanup>` can run more than once
for one extraction.
"""
##########################################################################
# 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, bidirectional: 'bool' = True,
analyse: 'bool' = False) -> 'None':
"""Initialise instance.
Arguments:
fout: output path
format: output format
byteorder: output file byte order
nanosecond: output nanosecond-resolution file flag
bidirectional: whether the two halves of a conversation are one flow.
:data:`True` -- the default -- keys a flow on the *pair* of
endpoints rather than on (source, destination), so a connection
is traced as the one thing it is; pass :data:`False` to trace
each direction as a flow of its own.
analyse: whether to reassemble each flow's application layer, so that
its ``packet`` can be read. Off by default: it buffers every
traced payload, a cost tracing does not otherwise pay.
"""
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
#: bool: Bidirectional tracing flag. If set to :data:`True`, both halves
#: of a conversation share one buffer entry, one label and one output
#: file; otherwise each direction is a flow of its own.
self._bidir = bidirectional
#: bool: Application-layer analysis flag. If set to :data:`True`, each
#: flow reassembles the payload it carries so that its ``packet`` can be
#: read; otherwise no payload is buffered and ``packet`` is
#: :data:`None`.
self._analyse = analyse
# dump I/O object
fio, ext = self.make_fout(fout, format)
#: ~typing.Type[Dumper]: Dumper class.
self._foutio = fio
#: Optional[str]: Output file extension.
self._fdpext = ext
logger.debug('%s flow tracing initialised (root=%s, format=%s, byteorder=%s, '
'nanosecond=%s, bidirectional=%s, analyse=%s)', self.name, fout,
format, byteorder, nanosecond, bidirectional, analyse)
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__ = []
[docs]
class TraceFlow(TraceFlowBase[_DT, _BT, _IT, _PT], Generic[_DT, _BT, _IT, _PT]):
"""Base flow tracing class.
Example:
Registration is opt-in. Pass keyword argument ``protocol`` at class
definition to register the flow tracing class under that protocol name:
.. code-block:: python
class MyProtocol(TraceFlow, protocol='my_protocol'):
...
Omit it and the subclass is *not* registered, which is how a class
that is not meant to be selectable by name declines:
.. code-block:: python
class MyMixin(TraceFlow): # not registered
...
Such a class can still be registered later, on demand:
.. code-block:: python
Extractor.register_traceflow('my_mixin', MyMixin)
Arguments:
fout: output path
format: output format
byteorder: output file byte order
nanosecond: output nanosecond-resolution file flag
bidirectional: trace both halves of a conversation as one flow
analyse: reassemble each flow's application layer
.. versionchanged:: 1.3.0
:attr:`name` and :attr:`protocol` are read from ``__protocol_name__``
and ``__protocol_type__``, where they were abstract properties each
flow tracing class implemented, and the built-in name is short:
``'TCP'`` rather than ``'Transmission Control Protocol'``. A subclass
registers itself with :class:`~pcapkit.foundation.extraction.Extractor`
at class definition.
.. versionchanged:: 1.5.0
Registration is opt-in, under the ``protocol`` class keyword; 1.4.x
registered every subclass, under its :attr:`name` when the keyword was
omitted (:issue:`514`). The ``bidirectional`` and ``analyse`` arguments
are added, and a flow covers both directions of a conversation unless
``bidirectional=False`` is passed.
"""
[docs]
def __init_subclass__(cls, /, protocol: 'Optional[str]' = None, *args: 'Any', **kwargs: 'Any') -> 'None':
"""Initialise subclass.
This method is to be used for registering the flow tracing class to
:class:`~pcapkit.foundation.extraction.Extractor` class.
Args:
protocol: Protocol name to register the subclass under, lowercased.
:data:`None` (the default) skips registration entirely.
*args: Arbitrary positional arguments.
**kwargs: Arbitrary keyword arguments.
Raises:
UnsupportedCall: If any unrecognised class keyword is given.
Registration is **opt-in**: the subclass is registered if and only if
``protocol`` is given. This is what lets a subclass decline registration
rather than having to inherit :class:`TraceFlowBase` to avoid it, and it
matches :meth:`EnumSchema.__init_subclass__
<pcapkit.protocols.schema.schema.EnumSchema.__init_subclass__>`, which
guards on its own ``code`` keyword the same way.
Note:
:attr:`__protocol_name__` is *not* an opt-in. It supplies the
:attr:`name <pcapkit.foundation.traceflow.traceflow.TraceFlow.name>`
the class reports, which it does whether or not the class is
registered; only the keyword decides registration.
See Also:
For more details, please refer to
:meth:`pcapkit.foundation.extraction.Extractor.register_traceflow`.
"""
# NOTE: the keyword here is ``protocol``, but ``Engine`` spells the same
# idea ``engine`` -- so guessing the other base's keyword, or ``name=``,
# is the expected mistake, not a careless one. Left in ``**kwargs`` it
# would be dropped by the bare ``super().__init_subclass__()`` below,
# leaving the class silently unregistered: no exception, no warning.
# See the sibling note in ``Engine.__init_subclass__``.
if args or kwargs:
unexpected = ', '.join([*map(repr, args), *sorted(kwargs)])
raise UnsupportedCall(f'{cls.__name__}: unexpected class keyword(s): {unexpected}')
if protocol is not None:
from pcapkit.foundation.extraction import \
Extractor # pylint: disable=import-outside-toplevel
Extractor.register_traceflow(protocol.lower(), cls)
return super().__init_subclass__()