Refactor TCP packet handler to use TcpOptions class for assembling op…

…tions

pytcp/protocols/tcp/assembler.py

@@ -24,6 +24,8 @@
 ############################################################################
 
 # pylint: disable = too-many-instance-attributes
+# pylint: disable = too-many-arguments
+# pylint: disable = too-many-locals
 
 """
 Module contains assembler support class for the TCP protocol.
@@ -43,14 +45,7 @@
 from pytcp.lib.tracker import Tracker
 from pytcp.protocols.tcp.base import Tcp
 from pytcp.protocols.tcp.header import TCP_HEADER_LEN, TcpHeader
-from pytcp.protocols.tcp.options import (
-    TcpOption,
-    TcpOptionEol,
-    TcpOptionMss,
-    TcpOptionNop,
-    TcpOptionSackPerm,
-    TcpOptionWscale,
-)
+from pytcp.protocols.tcp.options import TcpOptions
 
 
 class TcpAssembler(Tcp, ProtoAssembler):
@@ -77,7 +72,7 @@ def __init__(
         tcp__win: int = 0,
         tcp_cksum: int = 0,
         tcp__urg: int = 0,
-        tcp__options: list[TcpOption] | None = None,
+        tcp__options: TcpOptions | None = None,
         tcp__data: bytes | None = None,
         echo_tracker: Tracker | None = None,
     ) -> None:
@@ -96,11 +91,9 @@ def __init__(
 
         self._data = b"" if tcp__data is None else tcp__data
 
-        self._options: list[TcpOption] = (
-            [] if tcp__options is None else tcp__options
-        )
+        self._options = tcp__options or TcpOptions()
 
-        self._olen = sum(len(option) for option in self._options)
+        self._olen = len(self._options)
         self._hlen = TCP_HEADER_LEN + self._olen
         self._dlen = len(self._data)
         self._plen = self._hlen + self._dlen
@@ -145,103 +138,10 @@ def tracker(self) -> Tracker:
 
         return self._tracker
 
-    @property
-    def _raw_options(self) -> bytes:
-        """
-        Get packet options in raw format.
-        """
-
-        return b"".join(bytes(option) for option in self._options)
-
     def assemble(self, *, frame: memoryview, pshdr_sum: int) -> None:
         """
         Assemble packet into the frame.
         """
 
         struct.pack_into(f"{len(self)}s", frame, 0, bytes(self))
         struct.pack_into("! H", frame, 16, inet_cksum(frame, pshdr_sum))
-
-
-#
-# TCP options
-#
-
-
-class TcpOptionEolAssembler(TcpOptionEol):
-    """
-    TCP EOL option assembler.
-    """
-
-    def __init__(self) -> None:
-        """
-        Class constructor.
-        """
-
-
-class TcpOptionNopAssembler(TcpOptionNop):
-    """
-    TCP NOP option assembler.
-    """
-
-    def __init__(self) -> None:
-        """
-        Class constructor.
-        """
-
-
-class TcpOptionMssAssembler(TcpOptionMss):
-    """
-    TCP MSS option assembler.
-    """
-
-    def __init__(self, *, mss: int) -> None:
-        """
-        CLass constructor.
-        """
-
-        assert 0 <= mss <= 0xFFFF
-
-        self._mss = mss
-
-
-class TcpOptionWscaleAssembler(TcpOptionWscale):
-    """
-    TCP Wscale option assembler.
-    """
-
-    def __init__(self, *, wscale: int) -> None:
-        """
-        Class constructor.
-        """
-
-        assert 0 <= wscale <= 0xFF
-
-        self._wscale = wscale
-
-
-class TcpOptionSackPermAssembler(TcpOptionSackPerm):
-    """
-    TCP SackPerm option assembler.
-    """
-
-    def __init__(self) -> None:
-        """
-        Class constructor.
-        """
-
-
-class TcpOptionTimestampAssembler(TcpOptionWscale):
-    """
-    TCP Timestamp option assembler.
-    """
-
-    def __init__(self, *, tsval: int, tsecr: int) -> None:
-        """
-        Class constructor.
-        """
-
-        assert 0 <= tsval <= 0xFFFFFFFF
-        assert 0 <= tsecr <= 0xFFFFFFFF
-
-        self._tsval = tsval
-        self._tsecr = tsecr

pytcp/protocols/tcp/base.py

@@ -40,7 +40,7 @@
 from pytcp.protocols.ip4.header import Ip4Proto
 from pytcp.protocols.ip6.header import Ip6Next
 from pytcp.protocols.tcp.header import TcpHeader, TcpHeaderProperties
-from pytcp.protocols.tcp.options import TcpOption, TcpOptionsProperties
+from pytcp.protocols.tcp.options import TcpOptions, TcpOptionsProperties
 
 
 class Tcp(Proto, TcpHeaderProperties, TcpOptionsProperties):
@@ -52,7 +52,7 @@ class Tcp(Proto, TcpHeaderProperties, TcpOptionsProperties):
     __ip4_proto = Ip4Proto.TCP
 
     _header: TcpHeader
-    _options: list[TcpOption]
+    _options: TcpOptions
     _data: bytes
 
     _plen: int
@@ -66,20 +66,15 @@ def __str__(self) -> str:
         Get packet log string.
         """
 
-        log = (
+        return (
             f"TCP {self._header.sport} > {self._header.dport}, "
             f"{'N' if self._header.flag_ns else ''}{'C' if self._header.flag_cwr else ''}"
             f"{'E' if self._header.flag_ece else ''}{'U' if self._header.flag_urg else ''}"
             f"{'A' if self._header.flag_ack else ''}{'P' if self._header.flag_psh else ''}"
             f"{'R' if self._header.flag_rst else ''}{'S' if self._header.flag_syn else ''}"
             f"{'F' if self._header.flag_fin else ''}, seq {self._header.seq}, "
             f"ack {self._header.ack}, win {self._header.win}, dlen {len(self._data)}"
-        )
-
-        for option in self._options:
-            log += ", " + str(option)
-
-        return log
+        ) + f", opts [{self._options}]"
 
     @override
     def __repr__(self) -> str:
@@ -95,11 +90,7 @@ def __bytes__(self) -> bytes:
         Get the packet in raw form.
         """
 
-        return (
-            bytes(self._header)
-            + b"".join(bytes(option) for option in self._options)
-            + self._data
-        )
+        return bytes(self._header) + bytes(self._options) + self._data
 
     @property
     def ip6_next(self) -> Ip6Next:
@@ -118,7 +109,7 @@ def ip4_proto(self) -> Ip4Proto:
         return self.__ip4_proto
 
     @property
-    def options(self) -> list[TcpOption]:
+    def options(self) -> TcpOptions:
         """
         Get the options.
         """