Skip to content

Commit

Permalink
Refactor ICMPv6, TCP, IPv4 option handling to use from_bytes method.
Browse files Browse the repository at this point in the history
  • Loading branch information
@ccie18643
ccie18643 committed Jul 15, 2024
1 parent c0e73a4 commit b233aad
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 58 deletions.
20 changes: 9 additions & 11 deletions pytcp/protocols/icmp6/options__nd.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,9 +244,7 @@ def from_bytes(frame: bytes) -> Icmp6NdOptions:
)
case _:
options.append(
Icmp6NdOptionUnknown.from_frame(
frame=frame[option__ptr:]
)
Icmp6NdOptionUnknown.from_bytes(frame[option__ptr:])
)

option__ptr += options[-1].len
Expand Down Expand Up @@ -344,7 +342,7 @@ def slla(self) -> MacAddress:
@staticmethod
def from_bytes(bytes: bytes) -> Icmp6NdOptionSlla:
"""
Read the ICMPv6 ND option from frame.
Read the ICMPv6 ND option from bytes.
"""

return Icmp6NdOptionSlla(slla=MacAddress(bytes[2:8]))
Expand Down Expand Up @@ -407,7 +405,7 @@ def tlla(self) -> MacAddress:
@staticmethod
def from_bytes(bytes: bytes) -> Icmp6NdOptionTlla:
"""
Read the ICMPv6 ND option from frame.
Read the ICMPv6 ND option from bytes.
"""

return Icmp6NdOptionTlla(tlla=MacAddress(bytes[2:8]))
Expand Down Expand Up @@ -550,7 +548,7 @@ def prefix(self) -> Ip6Network:
@staticmethod
def from_bytes(bytes: bytes) -> Icmp6NdOptionPi:
"""
Read the ICMPv6 ND option from frame.
Read the ICMPv6 ND option from bytes.
"""

return Icmp6NdOptionPi(
Expand Down Expand Up @@ -629,13 +627,13 @@ def data(self) -> bytes:
return self._data

@staticmethod
def from_frame(frame: bytes) -> Icmp6NdOptionUnknown:
def from_bytes(bytes: bytes) -> Icmp6NdOptionUnknown:
"""
Read the ICMPv6 ND option from frame.
Read the ICMPv6 ND option from bytes.
"""

return Icmp6NdOptionUnknown(
type=Icmp6NdOptionType.from_frame(frame),
len=len(frame),
data=frame[2:],
type=Icmp6NdOptionType(bytes[0]),
len=bytes[0],
data=bytes[2:],
)
22 changes: 9 additions & 13 deletions pytcp/protocols/ip4/options.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,17 +127,13 @@ def from_frame(frame: bytes) -> Ip4Options:
while option__ptr < ip4__hlen:
match Ip4OptionType.from_frame(frame[option__ptr:]):
case Ip4OptionType.EOL:
options.append(
Ip4OptionEol.from_frame(frame=frame[option__ptr:])
)
options.append(Ip4OptionEol.from_bytes(frame[option__ptr:]))
break
case Ip4OptionType.NOP:
options.append(
Ip4OptionNop.from_frame(frame=frame[option__ptr:])
)
options.append(Ip4OptionNop.from_bytes(frame[option__ptr:]))
case _:
options.append(
Ip4OptionUnknown.from_frame(frame=frame[option__ptr:])
Ip4OptionUnknown.from_bytes(frame[option__ptr:])
)

option__ptr += options[-1].len
Expand Down Expand Up @@ -216,7 +212,7 @@ def __bytes__(self) -> bytes:
return struct.pack("! B", self._type.value)

@staticmethod
def from_frame(frame: bytes) -> Ip4OptionEol:
def from_bytes(_: bytes) -> Ip4OptionEol:
"""
Read the IPv4 EOL option from frame.
"""
Expand Down Expand Up @@ -262,7 +258,7 @@ def __bytes__(self) -> bytes:
return struct.pack("! B", self._type.value)

@staticmethod
def from_frame(frame: bytes) -> Ip4OptionNop:
def from_bytes(_: bytes) -> Ip4OptionNop:
"""
Read the IPv4 NOP option from frame.
"""
Expand Down Expand Up @@ -329,15 +325,15 @@ def data(self) -> bytes:
return self._data

@staticmethod
def from_frame(frame: bytes) -> Ip4OptionUnknown:
def from_bytes(bytes: bytes) -> Ip4OptionUnknown:
"""
Read the IPv4 unknown option from frame.
"""

return Ip4OptionUnknown(
type=Ip4OptionType.from_frame(frame=frame),
len=frame[1] << 3,
data=frame[2:],
type=Ip4OptionType(bytes[0]),
len=bytes[1] << 3,
data=bytes[2:],
)


Expand Down
62 changes: 28 additions & 34 deletions pytcp/protocols/tcp/options.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -184,33 +184,27 @@ def from_frame(frame: bytes) -> TcpOptions:
while option__ptr < tcp__hlen:
match TcpOptionType.from_frame(frame[option__ptr:]):
case TcpOptionType.EOL:
options.append(
TcpOptionEol.from_frame(frame=frame[option__ptr:])
)
options.append(TcpOptionEol.from_bytes(frame[option__ptr:]))
break
case TcpOptionType.NOP:
options.append(
TcpOptionNop.from_frame(frame=frame[option__ptr:])
)
options.append(TcpOptionNop.from_bytes(frame[option__ptr:]))
case TcpOptionType.MSS:
options.append(
TcpOptionMss.from_frame(frame=frame[option__ptr:])
)
options.append(TcpOptionMss.from_bytes(frame[option__ptr:]))
case TcpOptionType.WSCALE:
options.append(
TcpOptionWscale.from_frame(frame=frame[option__ptr:])
TcpOptionWscale.from_bytes(frame[option__ptr:])
)
case TcpOptionType.SACKPERM:
options.append(
TcpOptionSackPerm.from_frame(frame=frame[option__ptr:])
TcpOptionSackPerm.from_bytes(frame[option__ptr:])
)
case TcpOptionType.TIMESTAMP:
options.append(
TcpOptionTimestamp.from_frame(frame=frame[option__ptr:])
TcpOptionTimestamp.from_bytes(frame[option__ptr:])
)
case _:
options.append(
TcpOptionUnknown.from_frame(frame=frame[option__ptr:])
TcpOptionUnknown.from_bytes(frame[option__ptr:])
)

option__ptr += options[-1].len
Expand Down Expand Up @@ -289,9 +283,9 @@ def __bytes__(self) -> bytes:
return struct.pack("! B", self._type.value)

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionEol:
def from_bytes(_: bytes) -> TcpOptionEol:
"""
Create TCP EOL option from frame.
Create TCP EOL option from bytes.
"""

return TcpOptionEol()
Expand Down Expand Up @@ -335,9 +329,9 @@ def __bytes__(self) -> bytes:
return struct.pack("! B", self._type.value)

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionNop:
def from_bytes(_: bytes) -> TcpOptionNop:
"""
Create TCP NOP option from frame.
Create TCP NOP option from bytes.
"""

return TcpOptionNop()
Expand Down Expand Up @@ -399,12 +393,12 @@ def mss(self) -> int:
return self._mss

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionMss:
def from_bytes(bytes: bytes) -> TcpOptionMss:
"""
Create TCP MSS option from frame.
Create TCP MSS option from bytes.
"""

return TcpOptionMss(mss=struct.unpack_from("!H", frame, 2)[0])
return TcpOptionMss(mss=struct.unpack_from("!H", bytes, 2)[0])


class TcpOptionWscale(TcpOption):
Expand Down Expand Up @@ -463,12 +457,12 @@ def wscale(self) -> int:
return self._wscale

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionWscale:
def from_bytes(bytes: bytes) -> TcpOptionWscale:
"""
Create TCP Wscale option from frame.
Create TCP Wscale option from bytes.
"""

return TcpOptionWscale(wscale=frame[2])
return TcpOptionWscale(wscale=bytes[2])


class TcpOptionSackPerm(TcpOption):
Expand Down Expand Up @@ -513,9 +507,9 @@ def __bytes__(self) -> bytes:
)

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionSackPerm:
def from_bytes(_: bytes) -> TcpOptionSackPerm:
"""
Create TCP SackPerm option from frame.
Create TCP SackPerm option from bytes.
"""

return TcpOptionSackPerm()
Expand Down Expand Up @@ -591,14 +585,14 @@ def tsecr(self) -> int:
return self._tsecr

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionTimestamp:
def from_bytes(bytes: bytes) -> TcpOptionTimestamp:
"""
Create TCP Timestamp option from frame.
Create TCP Timestamp option from bytes.
"""

return TcpOptionTimestamp(
tsval=struct.unpack_from("!L", frame, 2)[0],
tsecr=struct.unpack_from("!L", frame, 6)[0],
tsval=struct.unpack_from("!L", bytes, 2)[0],
tsecr=struct.unpack_from("!L", bytes, 6)[0],
)


Expand Down Expand Up @@ -661,15 +655,15 @@ def data(self) -> bytes:
return self._data

@staticmethod
def from_frame(*, frame: bytes) -> TcpOptionUnknown:
def from_bytes(bytes: bytes) -> TcpOptionUnknown:
"""
Create TCP unknown option from frame.
Create TCP unknown option from bytes.
"""

return TcpOptionUnknown(
type=TcpOptionType.from_frame(frame=frame),
len=frame[1],
data=frame[2 : frame[1]],
type=TcpOptionType(bytes[0]),
len=bytes[1],
data=bytes[2 : bytes[1]],
)


Expand Down

0 comments on commit b233aad

Please sign in to comment.