Calculate checksums of known protocols for RAW payload.

Author: ccie18643

examples/icmp_echo_client.py

@@ -114,20 +114,38 @@ def _create_icmp4_message(cls, identifier: int, sequence: int) -> bytes:
         Create ICMPv4 Echo Request packet.
         """
 
-        def header(cksum: int = 0) -> bytes:
-            return struct.pack(
-                "!BBHHH",
-                ICMP4__ECHO_REQUEST__TYPE,
-                ICMP4__ECHO_REQUEST__CODE,
-                cksum,
-                identifier,
-                sequence,
-            )
+        header = struct.pack(
+            "!BBHHH",
+            ICMP4__ECHO_REQUEST__TYPE,
+            ICMP4__ECHO_REQUEST__CODE,
+            0,
+            identifier,
+            sequence,
+        )
+
+        payload = struct.pack("!d", time.time())
+
+        return header + payload
 
-        payload = struct.pack("!d", time.time())  # 8-byte timestamp
-        cksum = cls._checksum(header() + payload)
+    @classmethod
+    def _create_icmp6_message(cls, identifier: int, sequence: int) -> bytes:
+        """
+        Create ICMPv6 Echo Request packet.
+        """
 
-        return header(cksum) + payload
+        # Create the ICMPv6 header directly (without checksum)
+        header = struct.pack(
+            "!BBHHH",
+            ICMP6_ECHO_REQUEST_TYPE,
+            ICMP6_ECHO_REQUEST_CODE,
+            0,  # checksum (set to 0)
+            identifier,
+            sequence,
+        )
+
+        payload = struct.pack("!d", time.time())
+
+        return header + payload
 
     def _thread__client(self) -> None:
         """
@@ -142,6 +160,11 @@ def _thread__client(self) -> None:
                     self._local_ip_address.version,
                     self._remote_ip_address.version,
                 ):
+                    case 6, 6:
+                        icmp_message = self._create_icmp6_message(
+                            identifier=os.getpid() & 0xFFFF,
+                            sequence=self._message_count - message_count + 1,
+                        )
                     case 4, 4:
                         icmp_message = self._create_icmp4_message(
                             identifier=os.getpid() & 0xFFFF,

pytcp/protocols/ip4/ip4__base.py

@@ -44,12 +44,12 @@
 from pytcp.protocols.ip4.options.ip4_options import Ip4OptionsProperties
 from pytcp.protocols.tcp.tcp__assembler import TcpAssembler
 from pytcp.protocols.udp.udp__assembler import UdpAssembler
+from pytcp.protocols.raw.raw__assembler import RawAssembler
 
 if TYPE_CHECKING:
     from pytcp.protocols.icmp4.icmp4__assembler import Icmp4Assembler
     from pytcp.protocols.ip4.ip4__header import Ip4Header
     from pytcp.protocols.ip4.options.ip4_options import Ip4Options
-    from pytcp.protocols.raw.raw__assembler import RawAssembler
 
     Ip4Payload: TypeAlias = (
         Icmp4Assembler | TcpAssembler | UdpAssembler | RawAssembler
@@ -112,7 +112,9 @@ def __bytes__(self) -> bytes:
         )
         header_and_options[10:12] = inet_cksum(header_and_options).to_bytes(2)
 
-        if isinstance(self._payload, (TcpAssembler, UdpAssembler)):
+        if isinstance(
+            self._payload, (TcpAssembler, UdpAssembler, RawAssembler)
+        ):
             self._payload.pshdr_sum = self.pshdr_sum
 
         return bytes(header_and_options + bytes(self._payload))

pytcp/protocols/ip6/ip6__base.py

@@ -43,11 +43,11 @@
 from pytcp.protocols.ip6.ip6__header import Ip6HeaderProperties
 from pytcp.protocols.tcp.tcp__assembler import TcpAssembler
 from pytcp.protocols.udp.udp__assembler import UdpAssembler
+from pytcp.protocols.raw.raw__assembler import RawAssembler
 
 if TYPE_CHECKING:
     from pytcp.protocols.ip6.ip6__header import Ip6Header
     from pytcp.protocols.ip6_frag.ip6_frag__assembler import Ip6FragAssembler
-    from pytcp.protocols.raw.raw__assembler import RawAssembler
 
     Ip6Payload: TypeAlias = (
         Ip6FragAssembler
@@ -105,7 +105,8 @@ def __bytes__(self) -> bytes:
         """
 
         if isinstance(
-            self._payload, (TcpAssembler, UdpAssembler, Icmp6Assembler)
+            self._payload,
+            (TcpAssembler, UdpAssembler, Icmp6Assembler, RawAssembler),
         ):
             self._payload.pshdr_sum = self.pshdr_sum
 

pytcp/protocols/raw/raw__base.py

@@ -37,6 +37,7 @@
 
 from typing import override
 
+from pytcp.lib.inet_cksum import inet_cksum
 from pytcp.lib.proto import Proto
 from pytcp.protocols.enums import EtherType, IpProto
 
@@ -50,6 +51,8 @@ class Raw(Proto):
     _ether_type: EtherType
     _ip_proto: IpProto
 
+    pshdr_sum: int = 0
+
     @override
     def __len__(self) -> int:
         """
@@ -80,6 +83,39 @@ def __bytes__(self) -> bytes:
         Get the Raw packet as bytes.
         """
 
+        # Automatically calculate checksum if IpProto is ICMPv6 packet and checksum is not set.
+        if (
+            self._ip_proto == IpProto.ICMP6
+            and self._payload[2:4] == b"\x00\x00"
+        ):
+            _payload = bytearray(self._payload)
+            _payload[2:4] = inet_cksum(_payload, self.pshdr_sum).to_bytes(2)
+            return bytes(_payload)
+
+        # Automatically calculate checksum if IpProto is ICMPv4 packet and checksum is not set.
+        if (
+            self._ip_proto == IpProto.ICMP4
+            and self._payload[2:4] == b"\x00\x00"
+        ):
+            _payload = bytearray(self._payload)
+            _payload[2:4] = inet_cksum(_payload).to_bytes(2)
+            return bytes(_payload)
+
+        # Automatically calculate checksum if IpProto is ICMPv4 or ICMPv6 packet and checksum is not set.
+        if self._ip_proto == IpProto.UDP and self._payload[6:8] == b"\x00\x00":
+            _payload = bytearray(self._payload)
+            _payload[6:8] = inet_cksum(_payload, self.pshdr_sum).to_bytes(2)
+            return bytes(_payload)
+
+        # Automatically calculate checksum if IpProto is TCP packet and checksum is not set.
+        if (
+            self._ip_proto == IpProto.TCP
+            and self._payload[16:18] == b"\x00\x00"
+        ):
+            _payload = bytearray(self._payload)
+            _payload[16:18] = inet_cksum(_payload, self.pshdr_sum).to_bytes(2)
+            return bytes(_payload)
+
         return self._payload
 
     @property