Converted TCP to protocol header class

pytcp/protocols/tcp/fpa.py

@@ -45,6 +45,7 @@
 from pytcp.protocols.tcp.ps import (
     TCP_HEADER_LEN,
     Tcp,
+    TcpHeader,
     TcpOption,
     TcpOptionEol,
     TcpOptionMss,
@@ -76,6 +77,7 @@ def __init__(
         tcp__flag_syn: bool = False,
         tcp__flag_fin: bool = False,
         tcp__win: int = 0,
+        tcp_cksum: int = 0,
         tcp__urg: int = 0,
         tcp__options: list[TcpOption] | None = None,
         tcp__data: bytes | None = None,
@@ -94,31 +96,37 @@ def __init__(
 
         self._tracker: Tracker = Tracker(prefix="TX", echo_tracker=echo_tracker)
 
-        self._sport = tcp__sport
-        self._dport = tcp__dport
-        self._seq = tcp__seq
-        self._ack = tcp__ack
-        self._flag_ns = tcp__flag_ns
-        self._flag_cwr = tcp__flag_cwr
-        self._flag_ece = tcp__flag_ece
-        self._flag_urg = tcp__flag_urg
-        self._flag_ack = tcp__flag_ack
-        self._flag_psh = tcp__flag_psh
-        self._flag_rst = tcp__flag_rst
-        self._flag_syn = tcp__flag_syn
-        self._flag_fin = tcp__flag_fin
-        self._win = tcp__win
-        self._urg = tcp__urg
+        self._data = b"" if tcp__data is None else tcp__data
+
         self._options: list[TcpOption] = (
             [] if tcp__options is None else tcp__options
         )
-        self._data = b"" if tcp__data is None else tcp__data
 
         self._olen = sum(len(option) for option in self._options)
         self._hlen = TCP_HEADER_LEN + self._olen
         self._dlen = len(self._data)
         self._plen = self._hlen + self._dlen
 
+        self._header = TcpHeader(
+            sport=tcp__sport,
+            dport=tcp__dport,
+            seq=tcp__seq,
+            ack=tcp__ack,
+            hlen=self._hlen,
+            flag_ns=tcp__flag_ns,
+            flag_cwr=tcp__flag_cwr,
+            flag_ece=tcp__flag_ece,
+            flag_urg=tcp__flag_urg,
+            flag_ack=tcp__flag_ack,
+            flag_psh=tcp__flag_psh,
+            flag_rst=tcp__flag_rst,
+            flag_syn=tcp__flag_syn,
+            flag_fin=tcp__flag_fin,
+            win=tcp__win,
+            cksum=tcp_cksum,
+            urg=tcp__urg,
+        )
+
         assert self._hlen % 4 == 0, (
             f"TCP header len {self._hlen!r} is not multiplication of 4 bytes, "
             f"check options: {self._options!r}"

pytcp/protocols/tcp/fpp.py

@@ -49,6 +49,7 @@
     TCP_HEADER_LEN,
     TCP_OPTION_LEN__NOP,
     Tcp,
+    TcpHeader,
     TcpOption,
     TcpOptionEol,
     TcpOptionMss,
@@ -69,7 +70,7 @@ class TcpIntegrityError(PacketIntegrityError):
     Exception raised when TCP packet integrity check fails.
     """
 
-    def __init__(self, message: str):
+    def __init__(self, /, message: str):
         super().__init__("[TCP] " + message)
 
 
@@ -78,7 +79,7 @@ class TcpSanityError(PacketSanityError):
     Exception raised when TCP packet sanity check fails.
     """
 
-    def __init__(self, message: str):
+    def __init__(self, /, message: str):
         super().__init__("[TCP] " + message)
 
 
@@ -167,25 +168,10 @@ def _parse(self) -> None:
         Parse TCP packet.
         """
 
-        self._sport = struct.unpack("!H", self._frame[0:2])[0]
-        self._dport = struct.unpack("!H", self._frame[2:4])[0]
-        self._seq = struct.unpack("!L", self._frame[4:8])[0]
-        self._ack = struct.unpack("!L", self._frame[8:12])[0]
-        self._hlen = (self._frame[12] & 0b11110000) >> 2
-        self._flag_ns = bool(self._frame[12] & 0b00000001)
-        self._flag_cwr = bool(self._frame[13] & 0b10000000)
-        self._flag_ece = bool(self._frame[13] & 0b01000000)
-        self._flag_urg = bool(self._frame[13] & 0b00100000)
-        self._flag_ack = bool(self._frame[13] & 0b00010000)
-        self._flag_psh = bool(self._frame[13] & 0b00001000)
-        self._flag_rst = bool(self._frame[13] & 0b00000100)
-        self._flag_syn = bool(self._frame[13] & 0b00000010)
-        self._flag_fin = bool(self._frame[13] & 0b00000001)
-        self._win = struct.unpack("!H", self._frame[14:16])[0]
-        self._cksum = struct.unpack("!H", self._frame[16:18])[0]
-        self._urg = struct.unpack("!H", self._frame[18:20])[0]
+        self._header = TcpHeader.from_frame(self._frame)
 
         self._plen = self._ip__dlen
+        self._hlen = self._header.hlen
         self._olen = self._hlen - TCP_HEADER_LEN
         self._dlen = self._plen - self._hlen
 
@@ -236,42 +222,42 @@ def _validate_sanity(self) -> None:
         Check sanity of incoming packet after it has been parsed.
         """
 
-        if self._sport == 0:
+        if self._header.sport == 0:
             raise TcpSanityError(
                 "The 'sport' must be greater than 0.",
             )
 
-        if self._dport == 0:
+        if self._header.dport == 0:
             raise TcpSanityError(
                 "The 'dport' must be greater than  0.",
             )
 
-        if self._flag_syn and self._flag_fin:
+        if self._header.flag_syn and self._header.flag_fin:
             raise TcpSanityError(
                 "The 'flag_syn' and 'flag_fin' must not be set simultaneously.",
             )
 
-        if self._flag_syn and self._flag_rst:
+        if self._header.flag_syn and self._header.flag_rst:
             raise TcpSanityError(
                 "The 'flag_syn' and 'flag_rst' must not set simultaneously.",
             )
 
-        if self._flag_fin and self._flag_rst:
+        if self._header.flag_fin and self._header.flag_rst:
             raise TcpSanityError(
                 "The 'flag_fin' and 'flag_rst' must not be set simultaneously.",
             )
 
-        if self._flag_fin and not self._flag_ack:
+        if self._header.flag_fin and not self._header.flag_ack:
             raise TcpSanityError(
                 "The 'flag_ack' must be set when 'flag_fin' is set.",
             )
 
-        if self._ack and not self._flag_ack:
+        if self._header.ack and not self._header.flag_ack:
             raise TcpSanityError(
                 "The 'flag_ack' must be set when 'ack' is not 0.",
             )
 
-        if self._urg and not self._flag_urg:
+        if self._header.urg and not self._header.flag_urg:
             raise TcpSanityError(
                 "The 'flag_urg' must be set when 'urg' is not 0.",
             )