diff --git a/pytcp/lib/enum.py b/pytcp/lib/enum.py index 119eac99..f0b51f77 100755 --- a/pytcp/lib/enum.py +++ b/pytcp/lib/enum.py @@ -23,6 +23,7 @@ # # ############################################################################ +# pylint: disable = redefined-builtin """ Module contains the ProtoEnum class. @@ -96,16 +97,6 @@ def _extract( raise NotImplementedError - @staticmethod - def _from_bytes( - bytes: bytes, - ) -> int: - """ - Extract the enum value from the provided bytes. - """ - - raise NotImplementedError - @classmethod def from_frame( cls, @@ -122,16 +113,12 @@ def from_frame( return cls(value) @classmethod - def from_bytes( - cls, - /, - bytes: bytes, - ) -> Self: + def _from_bytes(cls, *, bytes: bytes, size: int) -> Self: """ - Create the enum object from the provided bytes. + Extract the enum value from the provided bytes. """ - if (value := cls._from_bytes(bytes)) not in cls: + if (value := int.from_bytes(bytes[:size])) not in cls: extend_enum(cls, f"UNKNOWN_{value}", value) return cls(value) @@ -151,3 +138,45 @@ def is_unknown(self) -> bool: """ return self.name.startswith("UNKNOWN_") + + +class ProtoEnumByte(ProtoEnum): + """ + Static enum used to represent protocol values stored in 8 bits. + """ + + @classmethod + def from_bytes(cls, /, bytes: bytes) -> Self: + """ + Extract the enum value from the provided bytes. + """ + + return cls._from_bytes(bytes=bytes, size=1) + + +class ProtoEnumWord(ProtoEnum): + """ + Static enum used to represent protocol values stored in 16 bits. + """ + + @classmethod + def from_bytes(cls, /, bytes: bytes) -> Self: + """ + Extract the enum value from the provided bytes. + """ + + return cls._from_bytes(bytes=bytes, size=2) + + +class ProtoEnumDWord(ProtoEnum): + """ + Static enum used to represent protocol values stored in 32 bits. + """ + + @classmethod + def from_bytes(cls, /, bytes: bytes) -> Self: + """ + Extract the enum value from the provided bytes. + """ + + return cls._from_bytes(bytes=bytes, size=4) diff --git a/pytcp/protocols/arp/header.py b/pytcp/protocols/arp/header.py index fb80b391..5ec7e293 100644 --- a/pytcp/protocols/arp/header.py +++ b/pytcp/protocols/arp/header.py @@ -40,9 +40,8 @@ import struct from abc import ABC from dataclasses import dataclass -from typing import override -from pytcp.lib.enum import ProtoEnum +from pytcp.lib.enum import ProtoEnumByte, ProtoEnumWord from pytcp.lib.ip4_address import Ip4Address from pytcp.lib.mac_address import MacAddress @@ -68,59 +67,39 @@ ARP_HEADER_LEN = 28 -class ArpHardwareType(ProtoEnum): +class ArpHardwareType(ProtoEnumWord): """ ARP hardware type. """ ETHERNET = 0x0001 - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:2]) - -class ArpProtocolType(ProtoEnum): +class ArpProtocolType(ProtoEnumWord): """ ARP protocol type. """ IP4 = 0x0800 - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:2]) - -class ArpHardwareLength(ProtoEnum): +class ArpHardwareLength(ProtoEnumByte): """ ARP hardware address length. """ ETHERNET = 6 - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:1]) - -class ArpProtocolLength(ProtoEnum): +class ArpProtocolLength(ProtoEnumByte): """ ARP protocol address length. """ IP4 = 4 - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:1]) - -class ArpOperation(ProtoEnum): +class ArpOperation(ProtoEnumWord): """ ARP operation. """ @@ -128,11 +107,6 @@ class ArpOperation(ProtoEnum): REQUEST = 1 REPLY = 2 - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:2]) - @dataclass class ArpHeader: diff --git a/pytcp/protocols/ethernet/header.py b/pytcp/protocols/ethernet/header.py index b5b9d5ba..ad55bcec 100644 --- a/pytcp/protocols/ethernet/header.py +++ b/pytcp/protocols/ethernet/header.py @@ -41,9 +41,8 @@ import struct from abc import ABC from dataclasses import dataclass -from typing import override -from pytcp.lib.enum import ProtoEnum +from pytcp.lib.enum import ProtoEnumWord from pytcp.lib.mac_address import MacAddress # Ethernet packet header. @@ -61,7 +60,7 @@ ETHERNET_HEADER_LEN = 14 -class EthernetType(ProtoEnum): +class EthernetType(ProtoEnumWord): """ Ethernet packet type enum. """ @@ -71,11 +70,6 @@ class EthernetType(ProtoEnum): IP6 = 0x86DD RAW = 0xFFFF - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:2]) - @dataclass class EthernetHeader: diff --git a/pytcp/protocols/icmp6/options__nd.py b/pytcp/protocols/icmp6/options__nd.py index 13835e92..3031b1e0 100644 --- a/pytcp/protocols/icmp6/options__nd.py +++ b/pytcp/protocols/icmp6/options__nd.py @@ -42,7 +42,7 @@ from dataclasses import dataclass from typing import override -from pytcp.lib.enum import ProtoEnum +from pytcp.lib.enum import ProtoEnumByte from pytcp.lib.ip6_address import Ip6Address, Ip6Mask, Ip6Network from pytcp.lib.mac_address import MacAddress from pytcp.lib.proto import Proto @@ -90,7 +90,7 @@ # TODO: Add ICMPv6 MTU Option (5). -class Icmp6NdOptionType(ProtoEnum): +class Icmp6NdOptionType(ProtoEnumByte): """ ICMPv6 Neighbor Discovery option types. """ @@ -99,10 +99,6 @@ class Icmp6NdOptionType(ProtoEnum): TLLA = 2 PI = 3 - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[0:1]) - ICMP6_ND_OPT_LEN__SLLA = 8 ICMP6_ND_OPT_LEN__TLLA = 8 diff --git a/pytcp/protocols/ip4/header.py b/pytcp/protocols/ip4/header.py index 7703b4b0..cc5550c6 100644 --- a/pytcp/protocols/ip4/header.py +++ b/pytcp/protocols/ip4/header.py @@ -42,7 +42,7 @@ from abc import ABC from dataclasses import dataclass -from pytcp.lib.enum import ProtoEnum +from pytcp.lib.enum import ProtoEnumByte from pytcp.lib.ip4_address import Ip4Address # IPv4 protocol header @@ -65,7 +65,7 @@ IP4_HEADER_LEN = 20 -class Ip4Proto(ProtoEnum): +class Ip4Proto(ProtoEnumByte): """ IPv4 protocol types. """ @@ -75,10 +75,6 @@ class Ip4Proto(ProtoEnum): UDP = 17 RAW = 255 - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[:1]) - @dataclass class Ip4Header: diff --git a/pytcp/protocols/ip4/options.py b/pytcp/protocols/ip4/options.py index e5ad02c7..6ebb6579 100644 --- a/pytcp/protocols/ip4/options.py +++ b/pytcp/protocols/ip4/options.py @@ -43,11 +43,11 @@ from abc import ABC from typing import override -from pytcp.lib.enum import ProtoEnum +from pytcp.lib.enum import ProtoEnumByte from pytcp.lib.proto import Proto -class Ip4OptionType(ProtoEnum): +class Ip4OptionType(ProtoEnumByte): """ IPv4 option types. """ @@ -55,11 +55,6 @@ class Ip4OptionType(ProtoEnum): EOL = 0 NOP = 1 - @override - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[0:1]) - IP4_OPTION_LEN__EOL = 1 IP4_OPTION_LEN__NOP = 1 diff --git a/pytcp/protocols/tcp/options.py b/pytcp/protocols/tcp/options.py index 115431fd..359ef9f1 100644 --- a/pytcp/protocols/tcp/options.py +++ b/pytcp/protocols/tcp/options.py @@ -42,13 +42,13 @@ from abc import ABC from typing import override -from pytcp.lib.enum import ProtoEnum +from pytcp.lib.enum import ProtoEnumByte from pytcp.lib.proto import Proto TCP_DEFAULT_MSS = 536 -class TcpOptionType(ProtoEnum): +class TcpOptionType(ProtoEnumByte): """ TCP option types. """ @@ -60,10 +60,6 @@ class TcpOptionType(ProtoEnum): SACKPERM = 4 TIMESTAMP = 8 - @staticmethod - def _from_bytes(bytes: bytes) -> int: - return int.from_bytes(bytes[0:1]) - TCP_OPTION_LEN__EOL = 1 TCP_OPTION_LEN__NOP = 1