fix(pcap,pcapng)!: a cut PCAP header or PCAP-NG block rebuilds as captured (#1470) (#1475)

Author: JarryShaw

pcapkit/protocols/data/misc/pcapng.py

@@ -69,6 +69,12 @@ class PCAPNG(Protocol):
     #: Block total length.
     length: 'int'
 
+    if TYPE_CHECKING:
+        #: The block's octets as captured, when the capture ended inside it;
+        #: present only then. A rebuild parses them again rather than writing
+        #: the block at its declared length (:issue:`1470`).
+        __truncated_raw__: 'bytes'
+
 
 @info_final
 class UnknownBlock(PCAPNG):

pcapkit/protocols/misc/pcap/header.py

@@ -35,14 +35,15 @@
 from pcapkit.protocols.data.misc.pcap.header import MagicNumber as Data_MagicNumber
 from pcapkit.protocols.protocol import ProtocolBase
 from pcapkit.protocols.schema.misc.pcap.header import Header as Schema_Header
+from pcapkit.protocols.schema.schema import keep_short_read, replay_short_read
 from pcapkit.utilities.exceptions import EndianError, FileError, UnsupportedCall
 
 if TYPE_CHECKING:
     from enum import IntEnum as StdlibEnum
     from typing import IO, Any, NoReturn, Optional, Type
 
     from aenum import IntEnum as AenumEnum
-    from typing_extensions import Literal
+    from typing_extensions import Literal, Self
 
 __all__ = ['Header']
 
@@ -256,6 +257,11 @@ def __post_init__(self, file: 'Optional[IO[bytes] | bytes]' = None,
         See Also:
             For construction argument, please refer to :meth:`make`.
 
+        Notes:
+            A header the data ends inside keeps that in its :attr:`info`, so
+            that :meth:`from_data` rebuilds only what was captured
+            (:issue:`1470`).
+
         """
         if file is None:
             _data = self.pack(**kwargs)
@@ -270,6 +276,7 @@ def __post_init__(self, file: 'Optional[IO[bytes] | bytes]' = None,
             self._file.name = file.name
         #: pcapkit.corekit.infoclass.Info: Parsed packet data.
         self._info = self.unpack(length, **kwargs)
+        keep_short_read(self)
 
     def __len__(self) -> 'Literal[24]':
         """Total length of corresponding protocol."""
@@ -279,6 +286,22 @@ def __length_hint__(self) -> 'Literal[24]':
         """Return an estimated length for the object."""
         return 24
 
+    @classmethod
+    def from_data(cls, data: 'Data_Header | dict[str, Any]', **kwargs: 'Any') -> 'Self':
+        """Create protocol instance from data.
+
+        Args:
+            data: Protocol data.
+            **kwargs: Construction keywords, as for :meth:`ProtocolBase.from_data
+                <pcapkit.protocols.protocol.ProtocolBase.from_data>`.
+
+        Returns:
+            Protocol instance, cut back to what was captured when the header
+            it was parsed from was cut short (:issue:`1470`).
+
+        """
+        return replay_short_read(super().from_data(data, **kwargs), data)
+
     @classmethod
     def __index__(cls) -> 'NoReturn':
         """Numeral registry index of the protocol.

pcapkit/protocols/misc/pcapng.py

@@ -15,6 +15,7 @@
 import base64
 import builtins
 import collections
+import collections.abc
 import datetime
 import decimal
 import enum
@@ -178,8 +179,9 @@
 from pcapkit.protocols.schema.misc.pcapng import ZigBeeNWKKey as Schema_ZigBeeNWKKey
 from pcapkit.protocols.schema.schema import Schema
 from pcapkit.utilities.compat import StrEnum, localcontext
-from pcapkit.utilities.exceptions import (EnumValueError, FormatError, ProtocolError, RegistryError,
-                                          StreamEOFError, UnsupportedCall, stacklevel)
+from pcapkit.utilities.exceptions import (BaseError, EnumValueError, FormatError, ProtocolError,
+                                          RegistryError, StreamEOFError, UnsupportedCall,
+                                          stacklevel)
 from pcapkit.utilities.warnings import (AttributeWarning, DeprecatedFormatWarning, ProtocolWarning,
                                         RegistryWarning, warn)
 
@@ -1264,7 +1266,15 @@ def read(self, length: 'Optional[int]' = None, *, _read: 'bool' = True,
                         getattr(self, meth_name, self._read_block_unknown))
         else:
             meth = name[0]
-        block = meth(schema.block, header=schema)
+        try:
+            block = meth(schema.block, header=schema)
+        except BaseError as error:
+            # NOTE: Fields past the end of a block the capture ended inside are
+            # read as zeros, and a value its parser refuses is an artefact of
+            # that, not of the block; it is kept as captured instead.
+            if not _read or self._captured(_seek_set) >= schema.block.length:
+                raise
+            block = self._read_block_cut(schema, _seek_set, error)
 
         # NOTE: An option or record area holding a malformed entry is kept as
         # the octets captured, as one :obj:`bytes` item (see
@@ -1314,6 +1324,14 @@ def read(self, length: 'Optional[int]' = None, *, _read: 'bool' = True,
                      stacklevel=stacklevel())
                 seek_cur = min(seek_cur, _seek_set + read)
 
+                # NOTE: The fields read past the end of the capture are zeros,
+                # and an option, record or packet the end cut through is only
+                # its first octets, so the block cannot be rebuilt from them at
+                # the length it declares. The octets captured are kept as
+                # ``__truncated_raw__`` instead, which :meth:`__post_init__` parses
+                # again in place of building the block (:issue:`1470`).
+                block.__update__([('__truncated_raw__', self._data)])
+
             # move backward to the beginning of next block
             self._file.seek(seek_cur, io.SEEK_SET)
 
@@ -1413,6 +1431,13 @@ def __post_init__(self, file: 'Optional[IO[bytes] | bytes]' = None, length: 'Opt
         #: collections.Counter: Counter for option types.
         self._opt = collections.Counter()  # type: Counter[Enum_OptionType]
 
+        # NOTE: A block the capture ended inside rebuilds as the octets that
+        # were captured, parsed again as they were the first time (see
+        # :meth:`read`), rather than at the length it declares (:issue:`1470`).
+        block = kwargs.get('block')
+        if file is None and isinstance(block, collections.abc.Mapping) and '__truncated_raw__' in block:
+            file = bytes(block['__truncated_raw__'])
+
         if file is None:
             _read = False
             #: bytes: Raw packet data.
@@ -1521,6 +1546,52 @@ def _check_block_floor(self, length: 'Optional[int]') -> 'None':
             raise StreamEOFError(f'PCAP-NG: block truncated: {length} octet(s) left, '
                                  'fewer than the 12 a block needs', quiet=True)
 
+    def _captured(self, offset: 'int') -> 'int':
+        """Count the octets left in the stream from ``offset`` on.
+
+        Args:
+            offset: File offset of the block.
+
+        Returns:
+            Octets from ``offset`` to the end of
+            :attr:`self._file <pcapkit.protocols.protocol.Protocol._file>`.
+
+        """
+        current = self._file.tell()
+        end = self._file.seek(0, io.SEEK_END)
+        self._file.seek(current, io.SEEK_SET)
+        return end - offset
+
+    def _read_block_cut(self, header: 'Schema_PCAPNG', offset: 'int',
+                        error: 'BaseError') -> 'Data_UnknownBlock':
+        """Read a PCAP-NG block the capture ended inside, whose parser refused it.
+
+        Args:
+            header: Parsed PCAP-NG header schema.
+            offset: File offset of the block.
+            error: What the block's parser raised.
+
+        Returns:
+            Parsed packet data, with the octets captured after the Block Total
+            Length as its body. :meth:`read` keeps the block's octets as
+            ``__truncated_raw__`` too, which a rebuild parses again, so that it
+            rebuilds as captured (:issue:`1470`).
+
+        """
+        warn(f'PCAP-NG: [Block {header.type}] block truncated; kept as captured: {error}',
+             ProtocolWarning, stacklevel=stacklevel())
+
+        current = self._file.tell()
+        self._file.seek(offset + 8, io.SEEK_SET)
+        body = self._read_fileng(header.block.length - 12)
+        self._file.seek(current, io.SEEK_SET)
+
+        return Data_UnknownBlock(
+            type=header.type,
+            length=header.block.length,
+            body=body,
+        )
+
     def _get_payload(self) -> 'bytes':
         """Get payload of :attr:`self.__header__ <pcapkit.protocols.protocol.Protocol.__header__>`.
 

pcapkit/protocols/schema/misc/pcapng.py

@@ -522,6 +522,25 @@ def pcapng_block_selector(packet: 'dict[str, Any]') -> 'Field':
     return SchemaField(length=max(packet['__length__'], 0), schema=schema)
 
 
+def _cut_last_line(schema: 'Schema', index: 'int', lines: 'list[str]') -> 'bool':
+    """Tell whether a key log line that did not parse is one the capture cut.
+
+    Args:
+        schema: Key log secrets schema, just unpacked.
+        index: Index of the line in ``lines``.
+        lines: Lines of the key log.
+
+    Returns:
+        Whether ``lines[index]`` is the last line of a key log the data ended
+        inside, which :meth:`Schema.unpack
+        <pcapkit.protocols.schema.schema.Schema.unpack>` records as
+        ``__short_read__``. Such a line is only the first octets of an entry,
+        and the log text, kept verbatim, still rebuilds it (:issue:`1470`).
+
+    """
+    return index == len(lines) - 1 and '__short_read__' in schema.__dict__
+
+
 def dsb_secrets_selector(packet: 'dict[str, Any]') -> 'Field':
     """Selector function for :attr:`DecryptionSecretsBlock.secrets_data` field.
 
@@ -2194,14 +2213,20 @@ def post_process(self, packet: 'dict[str, Any]') -> 'Schema':
         from pcapkit.protocols.misc.pcapng import TLSKeyLabel
 
         entries = collections.defaultdict(OrderedMultiDict)  # type: dict[TLSKeyLabel, OrderedMultiDict[bytes, bytes]]
-        for line in self.data.splitlines():
+        lines = self.data.splitlines()
+        for index, line in enumerate(lines):
             if not line or line.startswith('#'):
                 continue
 
-            label, random, secret = line.strip().split()
-            label_enum = TLSKeyLabel(label.upper())
-            entries[label_enum].add(bytes.fromhex(random),
-                                    bytes.fromhex(secret))
+            try:
+                label, random, secret = line.strip().split()
+                label_enum = TLSKeyLabel(label.upper())
+                entries[label_enum].add(bytes.fromhex(random),
+                                        bytes.fromhex(secret))
+            except ValueError as error:
+                if _cut_last_line(self, index, lines):
+                    break
+                raise FieldValueError(f'invalid TLS key log format: {line!r}') from error
 
         self.entries = entries
         return self
@@ -2233,16 +2258,24 @@ def post_process(self, packet: 'dict[str, Any]') -> 'Schema':
         from pcapkit.protocols.misc.pcapng import WireGuardKeyLabel
 
         entries = OrderedMultiDict()  # type: OrderedMultiDict[WireGuardKeyLabel, bytes]
-        for line in self.data.splitlines():
+        lines = self.data.splitlines()
+        for index, line in enumerate(lines):
             if not line or line.startswith('#'):
                 continue
 
             fields = line.strip().split()
-            if len(fields) != 3 or fields[1] != '=':
-                raise FieldValueError(f'invalid WireGuard key log format: {line!r}')
-            label, _, secret = fields
-            label_enum = WireGuardKeyLabel(label.upper())
-            entries.add(label_enum, base64.b64decode(secret))
+            try:
+                if len(fields) != 3 or fields[1] != '=':
+                    raise FieldValueError(f'invalid WireGuard key log format: {line!r}')
+                label, _, secret = fields
+                label_enum = WireGuardKeyLabel(label.upper())
+                entries.add(label_enum, base64.b64decode(secret))
+            except ValueError as error:
+                if _cut_last_line(self, index, lines):
+                    break
+                if isinstance(error, FieldValueError):
+                    raise
+                raise FieldValueError(f'invalid WireGuard key log format: {line!r}') from error
 
         self.entries = entries
         return self