fix(pcapng): rebuild a ZigBee key off its fixed width as captured (#1986)

docs/source/pcapkit/protocols/misc/pcapng.rst

@@ -581,6 +581,9 @@ shared with the PCAP frame header and documented in
 .. autofunction:: pcapkit.protocols.schema.misc.pcapng.bounded_area
 .. autofunction:: pcapkit.protocols.schema.misc.pcapng.captured_area
 .. autofunction:: pcapkit.protocols.schema.misc.pcapng.secrets_area
+.. autofunction:: pcapkit.protocols.schema.misc.pcapng.zigbee_key_has_trailing
+.. autofunction:: pcapkit.protocols.schema.misc.pcapng.zigbee_key_trailing
+.. autofunction:: pcapkit.protocols.schema.misc.pcapng.zigbee_key_held
 .. autofunction:: pcapkit.protocols.schema.misc.pcapng._spb_data_length
 
 Data Models

pcapkit/protocols/data/misc/pcapng.py

@@ -860,6 +860,11 @@ class ZigBeeNWKKey(DSBSecrets):
     pan_id: 'int'
 
     if TYPE_CHECKING:
+        #: Secrets Data octets as the block carries them, which a Secrets
+        #: Length may make shorter or longer than the key's fixed width. Set on
+        #: parsing only, and written back verbatim on rebuild (:issue:`1986`).
+        data: 'bytes'
+
         def __init__(self, nwk_key: 'bytes', pan_id: 'int') -> 'None': ...
 
 
@@ -875,6 +880,11 @@ class ZigBeeAPSKey(DSBSecrets):
     short_address: 'int'
 
     if TYPE_CHECKING:
+        #: Secrets Data octets as the block carries them, as for
+        #: :attr:`ZigBeeNWKKey.data`. Set on parsing only, and written back
+        #: verbatim on rebuild (:issue:`1986`).
+        data: 'bytes'
+
         def __init__(self, aps_key: 'bytes', pan_id: 'int', short_address: 'int') -> 'None': ...
 
 

pcapkit/protocols/misc/pcapng.py

@@ -208,6 +208,7 @@
 from pcapkit.protocols.schema.misc.pcapng import WireGuardKeyLog as Schema_WireGuardKeyLog
 from pcapkit.protocols.schema.misc.pcapng import ZigBeeAPSKey as Schema_ZigBeeAPSKey
 from pcapkit.protocols.schema.misc.pcapng import ZigBeeNWKKey as Schema_ZigBeeNWKKey
+from pcapkit.protocols.schema.misc.pcapng import zigbee_key_held
 from pcapkit.protocols.schema.schema import Schema
 from pcapkit.utilities.compat import StrEnum, localcontext
 from pcapkit.utilities.exceptions import (BaseError, EnumValueError, FormatError, ProtocolError,
@@ -4292,6 +4293,10 @@ def _read_secrets_zigbee_nwk(self, schema: 'Schema_ZigBeeNWKKey', *,
             nwk_key=schema.key,
             pan_id=schema.panid,
         )
+        # NOTE: the octets as the block holds them are kept as well, since a
+        # Secrets Length may declare fewer or more than the key's fixed width,
+        # and a rebuild writes them back verbatim (#1986).
+        secrets.__update__(data=bytes(schema))
         return secrets
 
     def _read_secrets_zigbee_aps(self, schema: 'Schema_ZigBeeAPSKey', *,
@@ -4342,6 +4347,8 @@ def _read_secrets_zigbee_aps(self, schema: 'Schema_ZigBeeAPSKey', *,
             pan_id=schema.panid,
             short_address=schema.addr_high << 16 | schema.addr_low,
         )
+        # NOTE: kept as for ``_read_secrets_zigbee_nwk`` (#1986).
+        secrets.__update__(data=bytes(schema))
         return secrets
 
     def _make_block_unknown(self, block: 'Optional[Data_UnknownBlock]' = None, *,
@@ -6716,7 +6723,12 @@ def _make_secrets_zigbee_nwk(self, type: 'Enum_SecretsType', secrets: 'Optional[
 
         Args:
             type: Secrets type.
-            secrets: Secrets data model.
+            secrets: Secrets data model. The Secrets Data octets it was parsed
+                from, if any, are written back verbatim, so that a key shorter
+                or longer than its fixed width rebuilds as captured
+                (:issue:`1986`); otherwise its fields are. On a parsed model
+                the octets win, so an edit to its fields alone is silently
+                ignored.
             nwk_key: NWK key.
             pan_id: PAN ID.
             **kwargs: Arbitrary keyword arguments.
@@ -6726,6 +6738,9 @@ def _make_secrets_zigbee_nwk(self, type: 'Enum_SecretsType', secrets: 'Optional[
 
         """
         if secrets is not None:
+            held = secrets.get('data')
+            if isinstance(held, bytes):
+                return zigbee_key_held(Schema_ZigBeeNWKKey, held)
             nwk_key = secrets.nwk_key
             pan_id = secrets.pan_id
 
@@ -6743,7 +6758,9 @@ def _make_secrets_zigbee_aps(self, type: 'Enum_SecretsType', secrets: 'Optional[
 
         Args:
             type: Secrets type.
-            secrets: Secrets data model.
+            secrets: Secrets data model. The Secrets Data octets it was parsed
+                from, if any, are written back verbatim, as for
+                :meth:`_make_secrets_zigbee_nwk`; otherwise its fields are.
             aps_key: APS key.
             pan_id: PAN ID.
             short_address: Short address.
@@ -6754,6 +6771,9 @@ def _make_secrets_zigbee_aps(self, type: 'Enum_SecretsType', secrets: 'Optional[
 
         """
         if secrets is not None:
+            held = secrets.get('data')
+            if isinstance(held, bytes):
+                return zigbee_key_held(Schema_ZigBeeAPSKey, held)
             aps_key = secrets.aps_key
             pan_id = secrets.pan_id
             short_address = secrets.short_address

pcapkit/protocols/schema/misc/pcapng.py

@@ -21,7 +21,8 @@
 from pcapkit.corekit.fields.collections import OptionField
 from pcapkit.corekit.fields.ipaddress import (IPv4AddressField, IPv4InterfaceField,
                                               IPv6AddressField, IPv6InterfaceField)
-from pcapkit.corekit.fields.misc import ForwardMatchField, PayloadField, SchemaField, SwitchField
+from pcapkit.corekit.fields.misc import (ConditionalField, ForwardMatchField, PayloadField,
+                                         SchemaField, SwitchField)
 from pcapkit.corekit.fields.numbers import (EnumField, Int32Field, Int64Field, NumberField,
                                             UInt8Field, UInt16Field, UInt32Field, UInt64Field)
 from pcapkit.corekit.fields.strings import (BitField, BytesField, DecodedString, PaddingField,
@@ -62,7 +63,7 @@
 
 if TYPE_CHECKING:
     from ipaddress import IPv4Address, IPv4Interface, IPv6Address, IPv6Interface
-    from typing import IO, Any, Callable, DefaultDict, Dict, Iterable, Optional, Type
+    from typing import IO, Any, Callable, DefaultDict, Dict, Iterable, Optional, Type, TypeVar
 
     from typing_extensions import Literal, Self
 
@@ -71,6 +72,8 @@
     from pcapkit.protocols.misc.pcapng import WireGuardKeyLabel
     from pcapkit.protocols.protocol import ProtocolBase
 
+    _ZigBeeKey = TypeVar('_ZigBeeKey', 'ZigBeeNWKKey', 'ZigBeeAPSKey')
+
 if SPHINX_TYPE_CHECKING:  # pragma: no cover
     from typing_extensions import TypedDict
 
@@ -611,6 +614,79 @@ def _secrets_cut(packet: 'dict[str, Any]') -> 'bool':
     return declared > secrets_area(block)
 
 
+def zigbee_key_has_trailing(packet: 'dict[str, Any]') -> 'bool':
+    """Tell whether a ZigBee key schema carries octets past its fixed fields.
+
+    Args:
+        packet: Unpacked data so far, or the data being packed.
+
+    Returns:
+        On unpacking, whether any of the span the block gives the key is left
+        once its fixed fields are read; on packing, whether the schema holds
+        such octets.
+
+    """
+    trailing = packet.get('trailing')
+    if isinstance(trailing, bytes):
+        return bool(trailing)
+    return packet['__length__'] > 0
+
+
+def zigbee_key_trailing(packet: 'dict[str, Any]') -> 'int':
+    """Size the octets a ZigBee key schema carries past its fixed fields.
+
+    Args:
+        packet: Unpacked data so far, or the data being packed.
+
+    Returns:
+        On unpacking, what is left of the span the block gives the key; on
+        packing, the octets the schema holds.
+
+    The ZigBee NWK and APS keys have a fixed width, but a block's Secrets
+    Length may declare more. The octets past the fixed fields are kept here,
+    rather than dropped, so that the key rebuilds as captured (:issue:`1986`).
+
+    """
+    trailing = packet.get('trailing')
+    if isinstance(trailing, bytes):
+        return len(trailing)
+    return max(packet['__length__'], 0)
+
+
+def zigbee_key_held(schema: 'Type[_ZigBeeKey]', data: 'bytes') -> '_ZigBeeKey':
+    """Rebuild a ZigBee key schema from the Secrets Data octets its block held.
+
+    Args:
+        schema: :class:`ZigBeeNWKKey` or :class:`ZigBeeAPSKey`.
+        data: The Secrets Data octets, as the block held them.
+
+    Returns:
+        A ``schema`` instance that packs back to exactly ``data``.
+
+    The fixed fields are read from ``data``, and any octets past them are kept
+    as the schema's ``trailing`` octets. A ``data`` short of the fixed width is
+    read as though zero-filled and recorded as a short read, so that packing
+    cuts it back to the octets held, as :meth:`Schema._cut_short_read
+    <pcapkit.protocols.schema.schema.Schema._cut_short_read>` describes. A key
+    shorter than its fixed width, or a Secrets Length longer than it, then
+    rebuilds byte for byte (:issue:`1986`).
+
+    """
+    fixed = [(name, field.length) for (name, field) in schema.__fields__.items()
+             if name != 'trailing']
+    width = sum(length for (_, length) in fixed)
+    padded = data.ljust(width, b'\x00')
+    value = cast('_ZigBeeKey', schema.unpack(padded))  # type: ignore[call-arg,misc]
+
+    offset = 0
+    for (name, length) in fixed:
+        if len(data) < offset + length:
+            value.__dict__['__short_read__'] = (name, len(data) - offset)
+            break
+        offset += length
+    return value
+
+
 def pcapng_block_selector(packet: 'dict[str, Any]') -> 'Field':
     """Selector function for :attr:`PCAPNG.block` field.
 
@@ -2481,13 +2557,20 @@ class ZigBeeNWKKey(DSBSecrets, code=Enum_SecretsType.ZigBee_NWK_Key):
     key: 'bytes' = BytesField(length=16)
     #: ZigBee PANID.
     panid: 'int' = UInt16Field(byteorder='little')
+    #: Octets past the PAN ID, if Secrets Length declares any (see
+    #: :func:`zigbee_key_trailing`).
+    trailing: 'Optional[bytes]' = ConditionalField(
+        BytesField(length=zigbee_key_trailing),
+        zigbee_key_has_trailing,
+    )
 
     # NOTE: The two zero octets after the PAN ID are the Decryption Secrets
     # Block's own padding, which Secrets Length does not count, so they belong
     # to :attr:`DecryptionSecretsBlock.padding_data` rather than to this schema.
 
     if TYPE_CHECKING:
-        def __init__(self, key: 'bytes', panid: 'int') -> 'None': ...
+        def __init__(self, key: 'bytes', panid: 'int',
+                     trailing: 'Optional[bytes]' = None) -> 'None': ...
 
 
 @schema_final
@@ -2502,12 +2585,19 @@ class ZigBeeAPSKey(DSBSecrets, code=Enum_SecretsType.ZigBee_APS_Key):
     addr_low: 'int' = UInt16Field(byteorder='little')
     #: High node short address.
     addr_high: 'int' = UInt16Field(byteorder='little')
+    #: Octets past the high node short address, if Secrets Length declares any
+    #: (see :func:`zigbee_key_trailing`).
+    trailing: 'Optional[bytes]' = ConditionalField(
+        BytesField(length=zigbee_key_trailing),
+        zigbee_key_has_trailing,
+    )
 
     # NOTE: As for :class:`ZigBeeNWKKey`, the trailing two zero octets are the
     # block's padding, not part of the secrets.
 
     if TYPE_CHECKING:
-        def __init__(self, key: 'bytes', panid: 'int', addr_low: 'int', addr_high: 'int') -> 'None': ...
+        def __init__(self, key: 'bytes', panid: 'int', addr_low: 'int', addr_high: 'int',
+                     trailing: 'Optional[bytes]' = None) -> 'None': ...
 
 
 class _DSB_Option(Option, ns='dsb'):

tests/protocols/misc/test_pcapng_dsb_secrets_overrun_1980_unit.py

@@ -91,9 +91,9 @@ def test_overrun_keeps_the_declared_length(self) -> None:
 
     def test_overrun_rebuilds(self) -> None:
         """The data model rebuilds the declared ``secrets_length``, as an EPB does its ``captured_len``."""
-        # NOTE: no ZigBee key here: one shorter than its fixed width rebuilds at
-        # full width whatever its secrets_length says, overrun or not.
-        for secrets_type in (TLS_KEY_LOG, UNKNOWN_SECRETS):
+        # NOTE: a ZigBee key shorter than its fixed width rebuilds as held too
+        # (#1986).
+        for secrets_type in (TLS_KEY_LOG, UNKNOWN_SECRETS, ZIGBEE_NWK_KEY):
             for secrets_length in (0xFFFF_FFF0, 5):
                 with self.subTest(secrets_type=hex(secrets_type), secrets_length=secrets_length):
                     octets = _dsb(secrets_type, secrets_length, b'abcd')
@@ -127,14 +127,19 @@ def assertValidBlock(self, octets: 'bytes', secrets_length: 'int') -> None:  # p
         self.assertEqual(struct.unpack('<I', octets[12:16])[0], secrets_length, octets.hex())
 
     def test_length_within_the_block_follows_the_secrets(self) -> None:
-        """A length that fits the block is no overrun, so the rebuild keeps its padding."""
+        """A length that fits the block is no overrun, so the rebuild keeps its padding.
+
+        The two octets past the ZigBee key's fixed width are secrets the length
+        declares, so they rebuild with it (#1986).
+        """
         key = bytes(range(16)) + b'\x34\x12'
         octets = _dsb(ZIGBEE_NWK_KEY, 20, key)
-        with warnings.catch_warnings():
-            warnings.simplefilter('ignore')
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter('always')
             rebuilt = self._rebuild(octets)
-        self.assertValidBlock(rebuilt, 18)
-        self.assertEqual(rebuilt[16:34], key)
+        self.assertFalse(any('runs past the block' in str(item.message) for item in caught))
+        self.assertValidBlock(rebuilt, 20)
+        self.assertEqual(rebuilt.hex(), octets.hex())
 
     def test_edited_secrets_take_their_own_length(self) -> None:
         from pcapkit.protocols.misc.pcapng import PCAPNG