Refactor Arp and Ethernet parsers to use from_bytes method for header…

… extraction

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)

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,71 +67,46 @@
 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.
     """
 
     REQUEST = 1
     REPLY = 2
 
-    @override
-    @staticmethod
-    def _from_bytes(bytes: bytes) -> int:
-        return int.from_bytes(bytes[:2])
-
 
 @dataclass
 class ArpHeader:

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:

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

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:

pytcp/protocols/ip4/options.py

@@ -43,23 +43,18 @@
 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.
     """
 
     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

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