From 0784938d3a48addb410c85241be8bbd54aad47c3 Mon Sep 17 00:00:00 2001
From: Sebastian Majewski <ccie18643@gmail.com>
Date: Mon, 15 Jul 2024 16:55:48 -0500
Subject: [PATCH] Refactor Arp and Ethernet parsers to use from_bytes method
 for header extraction

---
 pytcp/protocols/arp/header.py      |  6 +++++
 pytcp/protocols/ethernet/header.py |  2 ++
 pytcp/protocols/ip4/header.py      | 37 +++++++++++++++---------------
 pytcp/protocols/ip4/parser.py      |  2 +-
 4 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/pytcp/protocols/arp/header.py b/pytcp/protocols/arp/header.py
index baa8dcb3..fb80b391 100644
--- a/pytcp/protocols/arp/header.py
+++ b/pytcp/protocols/arp/header.py
@@ -40,6 +40,7 @@
 import struct
 from abc import ABC
 from dataclasses import dataclass
+from typing import override
 
 from pytcp.lib.enum import ProtoEnum
 from pytcp.lib.ip4_address import Ip4Address
@@ -74,6 +75,7 @@ class ArpHardwareType(ProtoEnum):
 
     ETHERNET = 0x0001
 
+    @override
     @staticmethod
     def _from_bytes(bytes: bytes) -> int:
         return int.from_bytes(bytes[:2])
@@ -86,6 +88,7 @@ class ArpProtocolType(ProtoEnum):
 
     IP4 = 0x0800
 
+    @override
     @staticmethod
     def _from_bytes(bytes: bytes) -> int:
         return int.from_bytes(bytes[:2])
@@ -98,6 +101,7 @@ class ArpHardwareLength(ProtoEnum):
 
     ETHERNET = 6
 
+    @override
     @staticmethod
     def _from_bytes(bytes: bytes) -> int:
         return int.from_bytes(bytes[:1])
@@ -110,6 +114,7 @@ class ArpProtocolLength(ProtoEnum):
 
     IP4 = 4
 
+    @override
     @staticmethod
     def _from_bytes(bytes: bytes) -> int:
         return int.from_bytes(bytes[:1])
@@ -123,6 +128,7 @@ class ArpOperation(ProtoEnum):
     REQUEST = 1
     REPLY = 2
 
+    @override
     @staticmethod
     def _from_bytes(bytes: bytes) -> int:
         return int.from_bytes(bytes[:2])
diff --git a/pytcp/protocols/ethernet/header.py b/pytcp/protocols/ethernet/header.py
index 9d96176c..b5b9d5ba 100644
--- a/pytcp/protocols/ethernet/header.py
+++ b/pytcp/protocols/ethernet/header.py
@@ -41,6 +41,7 @@
 import struct
 from abc import ABC
 from dataclasses import dataclass
+from typing import override
 
 from pytcp.lib.enum import ProtoEnum
 from pytcp.lib.mac_address import MacAddress
@@ -70,6 +71,7 @@ class EthernetType(ProtoEnum):
     IP6 = 0x86DD
     RAW = 0xFFFF
 
+    @override
     @staticmethod
     def _from_bytes(bytes: bytes) -> int:
         return int.from_bytes(bytes[:2])
diff --git a/pytcp/protocols/ip4/header.py b/pytcp/protocols/ip4/header.py
index d59a49d5..7703b4b0 100644
--- a/pytcp/protocols/ip4/header.py
+++ b/pytcp/protocols/ip4/header.py
@@ -24,6 +24,8 @@
 ############################################################################
 
 # pylint: disable = too-many-instance-attributes
+# pylint: disable = redefined-builtin
+
 
 """
 Module contains header support classes for the IPv4 protocol.
@@ -74,8 +76,8 @@ class Ip4Proto(ProtoEnum):
     RAW = 255
 
     @staticmethod
-    def _extract(frame: bytes) -> int:
-        return frame[9]
+    def _from_bytes(bytes: bytes) -> int:
+        return int.from_bytes(bytes[:1])
 
 
 @dataclass
@@ -126,27 +128,26 @@ def __bytes__(self) -> bytes:
         )
 
     @staticmethod
-    def from_frame(frame: bytes) -> Ip4Header:
+    def from_bytes(bytes: bytes) -> Ip4Header:
         """
         Populate the header from the raw frame.
         """
 
         return Ip4Header(
-            ver=frame[0] >> 4,
-            hlen=(frame[0] & 0b00001111) << 2,
-            dscp=(frame[1] & 0b11111100) >> 2,
-            ecn=frame[1] & 0b00000011,
-            plen=struct.unpack("!H", frame[2:4])[0],
-            id=struct.unpack("!H", frame[4:6])[0],
-            flag_df=bool(frame[6] & 0b01000000),
-            flag_mf=bool(frame[6] & 0b00100000),
-            offset=(struct.unpack("!H", frame[6:8])[0] & 0b0001111111111111)
-            << 3,
-            ttl=frame[8],
-            proto=Ip4Proto.from_frame(frame),
-            cksum=struct.unpack("!H", frame[10:12])[0],
-            src=Ip4Address(frame[12:16]),
-            dst=Ip4Address(frame[16:20]),
+            ver=bytes[0] >> 4,
+            hlen=(bytes[0] & 0b00001111) << 2,
+            dscp=(bytes[1] & 0b11111100) >> 2,
+            ecn=bytes[1] & 0b00000011,
+            plen=int.from_bytes(bytes[2:4]),
+            id=int.from_bytes(bytes[4:6]),
+            flag_df=bool(bytes[6] & 0b01000000),
+            flag_mf=bool(bytes[6] & 0b00100000),
+            offset=(int.from_bytes(bytes[6:8]) & 0b0001111111111111) << 3,
+            ttl=bytes[8],
+            proto=Ip4Proto.from_bytes(bytes[9:10]),
+            cksum=int.from_bytes(bytes[10:12]),
+            src=Ip4Address.from_bytes(bytes[12:16]),
+            dst=Ip4Address.from_bytes(bytes[16:20]),
         )
 
 
diff --git a/pytcp/protocols/ip4/parser.py b/pytcp/protocols/ip4/parser.py
index 37d5ca9a..72adcf3a 100755
--- a/pytcp/protocols/ip4/parser.py
+++ b/pytcp/protocols/ip4/parser.py
@@ -215,7 +215,7 @@ def _parse(self) -> None:
         Parse IPv4 packet.
         """
 
-        self._header = Ip4Header.from_frame(self._frame)
+        self._header = Ip4Header.from_bytes(self._frame)
         self._options = Ip4Options.from_bytes(
             self._frame[IP4_HEADER_LEN : self._header.hlen]
         )