From cdb6d6826ea7f0d89256f28309ecaacdf72e6f24 Mon Sep 17 00:00:00 2001 From: Sebastian Majewski Date: Sun, 14 Jul 2024 20:53:40 -0500 Subject: [PATCH] Refactor TCP packet handler to use TcpOptions class for assembling options --- pytcp/protocols/tcp/assembler.py | 112 +-------- pytcp/protocols/tcp/base.py | 21 +- pytcp/protocols/tcp/options.py | 284 +++++++++++++++++++++-- pytcp/protocols/tcp/packet_handler_tx.py | 22 +- pytcp/protocols/tcp/parser.py | 152 +----------- 5 files changed, 290 insertions(+), 301 deletions(-) diff --git a/pytcp/protocols/tcp/assembler.py b/pytcp/protocols/tcp/assembler.py index 4536605b..e7239858 100755 --- a/pytcp/protocols/tcp/assembler.py +++ b/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,14 +138,6 @@ 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. @@ -160,88 +145,3 @@ def assemble(self, *, frame: memoryview, pshdr_sum: int) -> None: 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 diff --git a/pytcp/protocols/tcp/base.py b/pytcp/protocols/tcp/base.py index 190381dc..6bbbc69f 100755 --- a/pytcp/protocols/tcp/base.py +++ b/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,7 +66,7 @@ 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 ''}" @@ -74,12 +74,7 @@ def __str__(self) -> str: 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. """ diff --git a/pytcp/protocols/tcp/options.py b/pytcp/protocols/tcp/options.py index adb0bcc2..493066fd 100644 --- a/pytcp/protocols/tcp/options.py +++ b/pytcp/protocols/tcp/options.py @@ -23,6 +23,10 @@ # # ############################################################################ +# pylint: disable = unused-argument +# pylint: disable = redefined-builtin + + """ Module contains options support classes for the TCP protccol. @@ -40,6 +44,7 @@ from pytcp.lib.enum import ProtoEnum from pytcp.lib.proto import Proto +from pytcp.protocols.tcp.header import TCP_HEADER_LEN TCP_DEFAULT_MSS = 536 @@ -69,6 +74,143 @@ def _extract(frame: bytes) -> int: TCP_OPTION_LEN__TIMESTAMP = 10 +class TcpOptions: + """ + TCP options. + """ + + _options: list[TcpOption] + + def __init__(self, *options: TcpOption) -> None: + """ + Initialize the TCP options. + """ + + self._options = list(options) + + def __len__(self) -> int: + """ + Get the options length. + """ + + return sum(len(option) for option in self._options) + + def __str__(self) -> str: + """ + Get the options log string. + """ + + return ", ".join(str(option) for option in self._options) + + def __repr__(self) -> str: + """ + Get the options representation string. + """ + + return f"TcpOptions(options={self._options!r})" + + def __bytes__(self) -> bytes: + """ + Get the options as bytes. + """ + + return b"".join(bytes(option) for option in self._options) + + @property + def mss(self) -> int: + """ + TCP option - Maximum Segment Size (2). + """ + + for option in self._options: + if isinstance(option, TcpOptionMss): + return option.mss + + return TCP_DEFAULT_MSS + + @property + def wscale(self) -> int | None: + """ + TCP option - Window Scale (3). + """ + + for option in self._options: + if isinstance(option, TcpOptionWscale): + return 1 << option.wscale + + return None + + @property + def sackperm(self) -> bool | None: + """ + TCP option - Sack Permit (4). + """ + + for option in self._options: + if isinstance(option, TcpOptionSackPerm): + return True + + return None + + @property + def timestamp(self) -> tuple[int, int] | None: + """ + TCP option - Timestamp (8). + """ + + for option in self._options: + if isinstance(option, TcpOptionTimestamp): + return option.tsval, option.tsecr + + return None + + @staticmethod + def from_frame(frame: bytes) -> TcpOptions: + """ + Read the TCP options from frame. + """ + + option__ptr = TCP_HEADER_LEN + tcp__hlen = (frame[12] >> 4) << 2 + + options: list[TcpOption] = [] + while option__ptr < tcp__hlen: + match TcpOptionType.from_frame(frame[option__ptr:]): + case TcpOptionType.EOL: + options.append( + TcpOptionEol.from_frame(frame=frame[option__ptr:]) + ) + break + case TcpOptionType.NOP: + options.append( + TcpOptionNop.from_frame(frame=frame[option__ptr:]) + ) + case TcpOptionType.MSS: + options.append( + TcpOptionMss.from_frame(frame=frame[option__ptr:]) + ) + case TcpOptionType.WSCALE: + options.append( + TcpOptionWscale.from_frame(frame=frame[option__ptr:]) + ) + case TcpOptionType.SACKPERM: + options.append( + TcpOptionSackPerm.from_frame(frame=frame[option__ptr:]) + ) + case TcpOptionType.TIMESTAMP: + options.append( + TcpOptionTimestamp.from_frame(frame=frame[option__ptr:]) + ) + case _: + options.append( + TcpOptionUnknown.from_frame(frame=frame[option__ptr:]) + ) + + option__ptr += options[-1].len + + return TcpOptions(*options) + + class TcpOption(Proto): """ Base class for TCP options. @@ -110,6 +252,11 @@ class TcpOptionEol(TcpOption): _type = TcpOptionType.EOL _len = TCP_OPTION_LEN__EOL + def __init__(self) -> None: + """ + Initialize the TCP EOL option. + """ + @override def __str__(self) -> str: """ @@ -134,6 +281,14 @@ def __bytes__(self) -> bytes: return struct.pack("! B", self._type.value) + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionEol: + """ + Create TCP EOL option from frame. + """ + + return TcpOptionEol() + class TcpOptionNop(TcpOption): """ @@ -143,6 +298,11 @@ class TcpOptionNop(TcpOption): _type = TcpOptionType.NOP _len = TCP_OPTION_LEN__NOP + def __init__(self) -> None: + """ + Initialize the TCP NOP option. + """ + @override def __str__(self) -> str: """ @@ -167,6 +327,14 @@ def __bytes__(self) -> bytes: return struct.pack("! B", self._type.value) + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionNop: + """ + Create TCP NOP option from frame. + """ + + return TcpOptionNop() + class TcpOptionMss(TcpOption): """ @@ -177,6 +345,15 @@ class TcpOptionMss(TcpOption): _len = TCP_OPTION_LEN__MSS _mss: int + def __init__(self, *, mss: int) -> None: + """ + Initialize the TCP MSS option. + """ + + assert 0 <= mss <= 0xFFFF + + self._mss = mss + @override def __str__(self) -> str: """ @@ -214,6 +391,14 @@ def mss(self) -> int: return self._mss + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionMss: + """ + Create TCP MSS option from frame. + """ + + return TcpOptionMss(mss=struct.unpack_from("!H", frame, 2)[0]) + class TcpOptionWscale(TcpOption): """ @@ -224,6 +409,15 @@ class TcpOptionWscale(TcpOption): _len = TCP_OPTION_LEN__WSCALE _wscale: int + def __init__(self, *, wscale: int) -> None: + """ + Initialize the TCP Wscale option. + """ + + assert 0 <= wscale <= 0xFF + + self._wscale = wscale + @override def __str__(self) -> str: """ @@ -261,6 +455,14 @@ def wscale(self) -> int: return self._wscale + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionWscale: + """ + Create TCP Wscale option from frame. + """ + + return TcpOptionWscale(wscale=frame[2]) + class TcpOptionSackPerm(TcpOption): """ @@ -270,6 +472,11 @@ class TcpOptionSackPerm(TcpOption): _type = TcpOptionType.SACKPERM _len = TCP_OPTION_LEN__SACKPERM + def __init__(self) -> None: + """ + Initialize the TCP SackPerm option. + """ + @override def __str__(self) -> str: """ @@ -298,6 +505,14 @@ def __bytes__(self) -> bytes: self._len, ) + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionSackPerm: + """ + Create TCP SackPerm option from frame. + """ + + return TcpOptionSackPerm() + class TcpOptionTimestamp(TcpOption): """ @@ -309,6 +524,17 @@ class TcpOptionTimestamp(TcpOption): _tsval: int _tsecr: int + def __init__(self, *, tsval: int, tsecr: int) -> None: + """ + Initialize the TCP Timestamp option. + """ + + assert 0 <= tsval <= 0xFFFFFFFF + assert 0 <= tsecr <= 0xFFFFFFFF + + self._tsval = tsval + self._tsecr = tsecr + @override def __str__(self) -> str: """ @@ -357,6 +583,17 @@ def tsecr(self) -> int: return self._tsecr + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionTimestamp: + """ + Create TCP Timestamp option from frame. + """ + + return TcpOptionTimestamp( + tsval=struct.unpack_from("!L", frame, 2)[0], + tsecr=struct.unpack_from("!L", frame, 6)[0], + ) + class TcpOptionUnknown(TcpOption): """ @@ -365,6 +602,15 @@ class TcpOptionUnknown(TcpOption): _data: bytes + def __init__(self, *, type: TcpOptionType, len: int, data: bytes) -> None: + """ + Initialize the TCP unknown option. + """ + + self._type = type + self._len = len + self._data = data + @override def __str__(self) -> str: """ @@ -407,13 +653,25 @@ def data(self) -> bytes: return self._data + @staticmethod + def from_frame(*, frame: bytes) -> TcpOptionUnknown: + """ + Create TCP unknown option from frame. + """ + + return TcpOptionUnknown( + type=TcpOptionType.from_frame(frame=frame), + len=frame[1], + data=frame[2 : frame[1]], + ) + class TcpOptionsProperties(ABC): """ Abstract class for TCP options properties. """ - _options: list[TcpOption] + _options: TcpOptions @property def mss(self) -> int: @@ -421,11 +679,7 @@ def mss(self) -> int: TCP option - Maximum Segment Size (2). """ - for option in self._options: - if isinstance(option, TcpOptionMss): - return option.mss - - return TCP_DEFAULT_MSS + return self._options.mss @property def wscale(self) -> int | None: @@ -433,11 +687,7 @@ def wscale(self) -> int | None: TCP option - Window Scale (3). """ - for option in self._options: - if isinstance(option, TcpOptionWscale): - return 1 << option.wscale - - return None + return self._options.wscale @property def sackperm(self) -> bool | None: @@ -445,11 +695,7 @@ def sackperm(self) -> bool | None: TCP option - Sack Permit (4). """ - for option in self._options: - if isinstance(option, TcpOptionSackPerm): - return True - - return None + return self._options.sackperm @property def timestamp(self) -> tuple[int, int] | None: @@ -457,8 +703,4 @@ def timestamp(self) -> tuple[int, int] | None: TCP option - Timestamp (8). """ - for option in self._options: - if isinstance(option, TcpOptionTimestamp): - return option.tsval, option.tsecr - - return None + return self._options.timestamp diff --git a/pytcp/protocols/tcp/packet_handler_tx.py b/pytcp/protocols/tcp/packet_handler_tx.py index c86cf704..7a2f127c 100755 --- a/pytcp/protocols/tcp/packet_handler_tx.py +++ b/pytcp/protocols/tcp/packet_handler_tx.py @@ -49,13 +49,13 @@ from pytcp.lib.logger import log from pytcp.lib.tracker import Tracker from pytcp.lib.tx_status import TxStatus -from pytcp.protocols.tcp.assembler import ( - TcpAssembler, - TcpOptionMssAssembler, - TcpOptionNopAssembler, - TcpOptionWscaleAssembler, +from pytcp.protocols.tcp.assembler import TcpAssembler +from pytcp.protocols.tcp.options import ( + TcpOptionMss, + TcpOptionNop, + TcpOptions, + TcpOptionWscale, ) -from pytcp.protocols.tcp.base import TcpOption class PacketHandlerTxTcp(ABC): @@ -126,17 +126,19 @@ def _phtx_tcp( self.packet_stats_tx.tcp__pre_assemble += 1 - options: list[TcpOption] = [] + options = TcpOptions() if tcp__mss: self.packet_stats_tx.tcp__opt_mss += 1 - options.append(TcpOptionMssAssembler(mss=tcp__mss)) + options = TcpOptions(TcpOptionMss(mss=tcp__mss)) if tcp__wscale: self.packet_stats_tx.tcp__opt_nop += 1 self.packet_stats_tx.tcp__opt_wscale += 1 - options.append(TcpOptionNopAssembler()) - options.append(TcpOptionWscaleAssembler(wscale=tcp__wscale)) + options = TcpOptions( + TcpOptionNop(), + TcpOptionWscale(wscale=tcp__wscale), + ) tcp_packet_tx = TcpAssembler( tcp__sport=tcp__sport, diff --git a/pytcp/protocols/tcp/parser.py b/pytcp/protocols/tcp/parser.py index 1da0a3c1..84e8817e 100755 --- a/pytcp/protocols/tcp/parser.py +++ b/pytcp/protocols/tcp/parser.py @@ -28,6 +28,7 @@ # pylint: disable = too-many-public-methods # pylint: disable = attribute-defined-outside-init + """ Module contains parser support class the for TCP protocol. @@ -39,7 +40,6 @@ from __future__ import annotations -import struct from typing import TYPE_CHECKING, override from pytcp.lib.errors import PacketIntegrityError, PacketSanityError @@ -49,15 +49,8 @@ from pytcp.protocols.tcp.header import TCP_HEADER_LEN, TcpHeader from pytcp.protocols.tcp.options import ( TCP_OPTION_LEN__NOP, - TcpOption, - TcpOptionEol, - TcpOptionMss, - TcpOptionNop, - TcpOptionSackPerm, - TcpOptionTimestamp, + TcpOptions, TcpOptionType, - TcpOptionUnknown, - TcpOptionWscale, ) if TYPE_CHECKING: @@ -174,44 +167,7 @@ def _parse(self) -> None: self._olen = self._hlen - TCP_HEADER_LEN self._dlen = self._plen - self._hlen - option__ptr = TCP_HEADER_LEN - self._options: list[TcpOption] = [] - - while option__ptr < self._hlen: - match TcpOptionType.from_frame(self._frame[option__ptr:]): - case TcpOptionType.EOL: - self._options.append( - TcpOptionEolParser(frame=self._frame[option__ptr:]) - ) - break - case TcpOptionType.NOP: - self._options.append( - TcpOptionNopParser(frame=self._frame[option__ptr:]) - ) - case TcpOptionType.MSS: - self._options.append( - TcpOptionMssParser(frame=self._frame[option__ptr:]) - ) - case TcpOptionType.WSCALE: - self._options.append( - TcpOptionWscaleParser(frame=self._frame[option__ptr:]) - ) - case TcpOptionType.SACKPERM: - self._options.append( - TcpOptionSackPermParser(frame=self._frame[option__ptr:]) - ) - case TcpOptionType.TIMESTAMP: - self._options.append( - TcpOptionTimestampParser( - frame=self._frame[option__ptr:] - ) - ) - case _: - self._options.append( - TcpOptionUnknownParser(frame=self._frame[option__ptr:]) - ) - - option__ptr += self._options[-1].len + self._options = TcpOptions.from_frame(frame=self._frame) self._data = self._frame[self._hlen : self._plen] @@ -260,105 +216,3 @@ def _validate_sanity(self) -> None: raise TcpSanityError( "The 'flag_urg' must be set when 'urg' is not 0.", ) - - -# -# TCP options -# - - -class TcpOptionEolParser(TcpOptionEol): - """ - TCP EOL option parser. - """ - - def __init__(self, *, frame: bytes) -> None: - """ - Class constructor. - """ - - -class TcpOptionNopParser(TcpOptionNop): - """ - TCP NOP option parser. - """ - - def __init__(self, *, frame: bytes) -> None: - """ - Class constructor. - """ - - -class TcpOptionMssParser(TcpOptionMss): - """ - TCP MSS option parser. - """ - - def __init__(self, *, frame: bytes) -> None: - """ - Class constructor. - """ - - self._type = TcpOptionType.from_frame(frame) - self._len = frame[1] - self._mss = struct.unpack_from("!H", frame, 2)[0] - - -class TcpOptionWscaleParser(TcpOptionWscale): - """ - TCP Wscale option parser. - """ - - def __init__(self, *, frame: bytes) -> None: - """ - Class constructor. - """ - - self._type = TcpOptionType.from_frame(frame) - self._len = frame[1] - self._wscale = frame[2] - - -class TcpOptionSackPermParser(TcpOptionSackPerm): - """ - TCP SackPerm option parser. - """ - - def __init__(self, *, frame: bytes) -> None: - """ - Class constructor. - """ - - self._type = TcpOptionType.from_frame(frame) - self._len = frame[1] - - -class TcpOptionTimestampParser(TcpOptionTimestamp): - """ - TCP Timestamp option parser. - """ - - def __init__(self, *, frame: bytes) -> None: - """ - Class constructor. - """ - - self._type = TcpOptionType.from_frame(frame) - self._len = frame[1] - self._tsval: int = struct.unpack_from("!L", frame, 2)[0] - self._tsecr: int = struct.unpack_from("!L", frame, 6)[0] - - -class TcpOptionUnknownParser(TcpOptionUnknown): - """ - IPv4 unknown option parser. - """ - - def __init__(self, frame: bytes) -> None: - """ - Class constructor. - """ - - self._type = TcpOptionType.from_frame(frame) - self._len = frame[1] - self._data = frame[2 : self._len]