diff --git a/scapy/packet.py b/scapy/packet.py index 3fb2eac6d92..27e8fbc796e 100644 --- a/scapy/packet.py +++ b/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 diff --git a/test/regression.uts b/test/regression.uts index 4beb41787ac..26d3d26fb6b 100644 --- a/test/regression.uts +++ b/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