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

Author: guedou

scapy/packet.py

@@ -79,6 +79,28 @@
 _T = TypeVar("_T", Dict[str, Any], Optional[Dict[str, Any]])
 
 
+def _rebuild_pkt(
+    cls,  # type: Type[Packet]
+    fields,  # type: Dict[str, Any]
+    payload,  # type: Optional[Packet]
+    metadata,  # 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']
+    return pkt
+
+
 class Packet(
     BasePacket,
     _CanvasDumpExtended,
@@ -214,15 +236,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]
@@ -244,27 +257,32 @@ def comment(self, value):
             self.comments = None
 
     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,
+        }
+        return (
+            _rebuild_pkt,
+            (self.__class__, fields, payload, metadata),
+        )
 
     def __deepcopy__(self,
                      memo,  # type: Any

test/regression.uts

@@ -616,6 +616,43 @@ 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)
+
 = Usage test
 
 from scapy.main import _usage