Pickle - reconstructs the packet from field values, payload and metadata (#4919)

* Pickle - reconstructs the packet from field valies, payload and metadata

* Pickle - support extra slots

Author: guedou

scapy/packet.py

@@ -214,15 +214,6 @@ def __init__(self,
         else:
             self.post_transforms = [post_transform]
 
-    _PickleType = Tuple[
-        Union[EDecimal, float],
-        Optional[Union[EDecimal, float, None]],
-        Optional[int],
-        Optional[_GlobInterfaceType],
-        Optional[int],
-        Optional[bytes],
-    ]
-
     @property
     def comment(self):
         # type: () -> Optional[bytes]
@@ -243,28 +234,64 @@ def comment(self, value):
         else:
             self.comments = None
 
+    @classmethod
+    def _rebuild_pkt(
+        cls,  # type: Type[Packet]
+        fields,  # type: Dict[str, Any]
+        payload,  # type: Optional[Packet]
+        metadata,  # type: Dict[str, Any]
+        extra_slots={},  # type: Dict[str, Any]
+    ):
+        # type: (...) -> Packet
+        """Helper for unpickling Packet instances via field values."""
+        # Create the instance using the field values
+        pkt = cls(**fields)
+        if payload is not None:
+            pkt.add_payload(payload)
+        # Restore metadata
+        pkt.time = metadata['time']
+        pkt.sent_time = metadata['sent_time']
+        pkt.direction = metadata['direction']
+        pkt.sniffed_on = metadata['sniffed_on']
+        pkt.wirelen = metadata['wirelen']
+        pkt.comments = metadata['comments']
+        # Restore any extra __slots__ defined by subclasses
+        for attr, value in extra_slots.items():
+            setattr(pkt, attr, value)
+        return pkt
+
     def __reduce__(self):
-        # type: () -> Tuple[Type[Packet], Tuple[bytes], Packet._PickleType]
-        """Used by pickling methods"""
-        return (self.__class__, (self.build(),), (
-            self.time,
-            self.sent_time,
-            self.direction,
-            self.sniffed_on,
-            self.wirelen,
-            self.comment
-        ))
-
-    def __setstate__(self, state):
-        # type: (Packet._PickleType) -> Packet
-        """Rebuild state using pickable methods"""
-        self.time = state[0]
-        self.sent_time = state[1]
-        self.direction = state[2]
-        self.sniffed_on = state[3]
-        self.wirelen = state[4]
-        self.comment = state[5]
-        return self
+        # type: () -> Tuple[Any, ...]
+        """Used by pickling methods.
+
+        Reconstructs the packet from field values, payload, and metadata.
+        """
+        # Store field values for unpickling
+        fields = {}
+        for f in self.fields_desc:
+            if f.name in self.fields:
+                fields[f.name] = self.fields[f.name]
+        payload = self.payload  # type: Optional[Packet]
+        if isinstance(payload, NoPayload):
+            payload = None
+        # Store metadata for unpickling
+        metadata = {
+            'time': self.time,
+            'sent_time': self.sent_time,
+            'direction': self.direction,
+            'sniffed_on': self.sniffed_on,
+            'wirelen': self.wirelen,
+            'comments': self.comments,
+        }
+        # Collect any extra __slots__ defined by subclasses
+        extra_slots = {}
+        for attr in type(self).__all_slots__ - set(Packet.__slots__):
+            if hasattr(self, attr):
+                extra_slots[attr] = getattr(self, attr)
+        return (
+            type(self)._rebuild_pkt,
+            (fields, payload, metadata, extra_slots),
+        )
 
     def __deepcopy__(self,
                      memo,  # type: Any

test/regression.uts

@@ -616,6 +616,74 @@ c = pickle.loads(b)
 assert c[IP].dst == "192.168.0.1"
 assert raw(c) == raw(a)
 
+= Pickle preserves field values, payload and metadata
+
+import pickle
+
+p = IP(src='1.2.3.4', dst='5.6.7.8')/TCP(sport=1234, dport=80, flags='S')
+p.time = 12345.0
+p.sent_time = 12346.0
+p.direction = 1
+p.sniffed_on = 'eth0'
+p.wirelen = 100
+p.comment = b'test comment'
+p2 = pickle.loads(pickle.dumps(p))
+assert p2[IP].src == '1.2.3.4'
+assert p2[IP].dst == '5.6.7.8'
+assert p2[TCP].sport == 1234
+assert p2[TCP].dport == 80
+assert p2[TCP].flags == 'S'
+assert p2[IP].len is None
+assert p2[IP].chksum is None
+assert p2[TCP].chksum is None
+assert p2.time == 12345.0
+assert p2.sent_time == 12346.0
+assert p2.direction == 1
+assert p2.sniffed_on == 'eth0'
+assert p2.wirelen == 100
+assert p2.comment == b'test comment'
+assert raw(p2) == raw(p)
+
+= Pickle a bare packet without payload
+
+import pickle
+
+p = IP(src='10.0.0.1')
+p2 = pickle.loads(pickle.dumps(p))
+assert p2.src == '10.0.0.1'
+assert raw(p2) == raw(p)
+
+= Pickle preserves custom __slots__ from subclasses
+
+import pickle
+import scapy.packet as _pkt_mod
+
+class _PickleTestPacket(Packet):
+    __slots__ = ["custom_id", "custom_tag"]
+    name = "PickleTestPacket"
+    fields_desc = [ByteField("val", 0)]
+
+# Make the class discoverable by pickle
+_pkt_mod._PickleTestPacket = _PickleTestPacket
+_PickleTestPacket.__module__ = 'scapy.packet'
+
+p = _PickleTestPacket(val=42)
+p.custom_id = 0x123
+p.custom_tag = "hello"
+p2 = pickle.loads(pickle.dumps(p))
+assert p2.val == 42
+assert p2.custom_id == 0x123
+assert p2.custom_tag == "hello"
+
+# Slots not explicitly set are not serialized
+p3 = _PickleTestPacket(val=7)
+assert not hasattr(p3, 'custom_id')
+p4 = pickle.loads(pickle.dumps(p3))
+assert p4.val == 7
+assert not hasattr(p4, 'custom_id')
+
+del _pkt_mod._PickleTestPacket
+
 = Usage test
 
 from scapy.main import _usage