diff --git a/pydivert/packet/header.py b/pydivert/packet/header.py index bdeb21a..ab67900 100644 --- a/pydivert/packet/header.py +++ b/pydivert/packet/header.py @@ -48,9 +48,9 @@ def raw(self) -> memoryview: @raw.setter def raw(self, val: bytes | bytearray | memoryview) -> None: if len(val) == len(self.raw): - self.raw[:] = val # pragma: no cover + self.raw[:] = val else: - self._packet.raw = memoryview(bytearray(self._packet.raw[: self._start].tobytes() + val)) + self._packet.raw = memoryview(bytearray(self._packet.raw[: self._start].tobytes() + bytes(val))) if self._packet.ip: self._packet.ip.packet_len = len(self._packet.raw) self._packet._invalidate_checksums() @@ -81,9 +81,9 @@ def payload(self) -> bytes: @payload.setter def payload(self, val: bytes | bytearray | memoryview) -> None: if len(val) == len(self.raw) - self.header_len: - self.raw[self.header_len :] = val # pragma: no cover + self.raw[self.header_len :] = val else: - self.raw = self.raw[: self.header_len].tobytes() + val + self.raw = self.raw[: self.header_len].tobytes() + bytes(val) class PortMixin(RawProtocol): @@ -92,19 +92,19 @@ def src_port(self) -> int: """ The source port. """ - return struct.unpack_from("!H", self.raw, 0)[0] # pragma: no cover + return struct.unpack_from("!H", self.raw, 0)[0] @src_port.setter def src_port(self, val: int) -> None: - self.raw[0:2] = struct.pack("!H", val) # pragma: no cover + self.raw[0:2] = struct.pack("!H", val) @property def dst_port(self) -> int: """ The destination port. """ - return struct.unpack_from("!H", self.raw, 2)[0] # pragma: no cover + return struct.unpack_from("!H", self.raw, 2)[0] @dst_port.setter def dst_port(self, val: int) -> None: - self.raw[2:4] = struct.pack("!H", val) # pragma: no cover + self.raw[2:4] = struct.pack("!H", val) diff --git a/pydivert/tests/test_packet.py b/pydivert/tests/test_packet.py index 3840672..7c11ea7 100644 --- a/pydivert/tests/test_packet.py +++ b/pydivert/tests/test_packet.py @@ -232,3 +232,57 @@ def test_ipv4_hypo(src, dst): p.dst_addr = dst assert p.src_addr == src assert p.dst_addr == dst + + +def test_header_raw_modification(): + p = create_base_ipv4_tcp() + assert p.ipv4 is not None + original_len = len(p.ipv4.raw) + + # same length + new_raw = bytearray([0x12] * original_len) + p.ipv4.raw = new_raw + assert bytes(p.ipv4.raw) == new_raw + + # different length + p = create_base_ipv4_tcp() + assert p.ipv4 is not None + new_raw = bytearray([0x45, 0x00, 0x00, 0x29] + [0x34] * (original_len - 4 + 1)) + p.ipv4.raw = new_raw + assert p.ipv4.packet_len == len(p.raw) + assert getattr(p, "_ip_checksum", False) is False + + # PayloadMixin and PortMixin + + from pydivert.packet.header import PayloadMixin, PortMixin + + class DummyProtocol(PortMixin, PayloadMixin): + def __init__(self): + self._raw = bytearray([0] * 20) + + @property + def raw(self): + return memoryview(self._raw) + + @raw.setter + def raw(self, val): + self._raw = bytearray(val) + + @property + def header_len(self): + return 8 + + proto = DummyProtocol() + assert proto.src_port == 0 + proto.src_port = 80 + assert proto.src_port == 80 + assert proto.dst_port == 0 + proto.dst_port = 8080 + assert proto.dst_port == 8080 + assert len(proto.payload) == 12 + proto.payload = bytearray([1] * 12) + assert len(proto.payload) == 12 + assert proto.payload[0] == 1 + proto.payload = bytearray([2] * 10) + assert len(proto.payload) == 10 + assert proto.payload[0] == 2