Refactor and improve code quality

optimus/dns/dns_packet.py

@@ -1,8 +1,7 @@
 from enum import Enum
 from typing import List, Optional, Union
 
-from optimus.dns.dns_records import (AAAA, NS, A, Record, RecordClass,
-                                     RecordType)
+from optimus.dns.dns_records import AAAA, NS, A, Record, RecordClass, RecordType
 
 
 class ResponseCode(Enum):  # 4 bits
@@ -49,9 +48,7 @@ def to_bin(self) -> bytearray:
         labels = self.name.split(".")
         for label in labels:
             # Write label's length
-            dns_question_bin.append(
-                len(label)
-            )  # TODO: Add check to ensure that label length is <=63
+            dns_question_bin.append(len(label))  # TODO: Add check to ensure that label length is <=63
             for ch in label:
                 data = ord(ch) if ch != "." else 0
                 dns_question_bin.append(data)

optimus/dns/dns_records.py

@@ -40,9 +40,7 @@ class Record:
     ttl: int  # 4 bytes : Time-to-Live in seconds
     length: int  # 2 bytes : Length of content in a concrete record
 
-    def __init__(
-        self, name: str, rtype: RecordType, rclass: RecordClass, ttl: int, length: int
-    ) -> None:
+    def __init__(self, name: str, rtype: RecordType, rclass: RecordClass, ttl: int, length: int) -> None:
         self.name = name
         self.rtype = rtype
         self.rec_class = rclass
@@ -54,9 +52,7 @@ def to_bin(self) -> bytearray:
         labels = self.name.split(".")
         for label in labels:
             # Write label's length
-            dns_record_bin.append(
-                len(label)
-            )  # TODO: Add check to ensure that label length is <=63
+            dns_record_bin.append(len(label))  # TODO: Add check to ensure that label length is <=63
             for ch in label:
                 data = ord(ch) if ch != "." else 0
                 dns_record_bin.append(data)
@@ -235,12 +231,8 @@ def __repr__(self) -> str:
 # Record Type MX, representing the host of the mail server for a domain
 class MX(Record):
     # Lower pref value => High Priority
-    preference: (
-        int  # 2 bytes : Specifies the preference given to this record among others
-    )
-    exchange: (
-        str  # Domain name which specifies a host willing to act as a mail exchange
-    )
+    preference: int  # 2 bytes : Specifies the preference given to this record among others
+    exchange: str  # Domain name which specifies a host willing to act as a mail exchange
 
     def __init__(
         self,

optimus/dns/handlers.py

@@ -5,8 +5,7 @@
 
 from optimus.dns.dns_packet import DNSHeader, DNSPacket, Question, ResponseCode
 from optimus.dns.dns_parser import DNSParser
-from optimus.dns.dns_records import (AAAA, NS, A, Record, RecordClass,
-                                     RecordType)
+from optimus.dns.dns_records import AAAA, NS, A, Record, RecordClass, RecordType
 from optimus.logging_config.logger import log
 from optimus.networking.udp_utils import query_server_over_udp
 
@@ -46,10 +45,7 @@ def __perform_recursive_lookup(qpacket: DNSPacket) -> DNSPacket:
             ResponseCode.UNKNOWN.value,
         ]:
             return response_packet
-        if (
-            response_code.value == ResponseCode.NOERROR.value
-            and len(response_packet.answers) > 0
-        ):
+        if response_code.value == ResponseCode.NOERROR.value and len(response_packet.answers) > 0:
             return response_packet
         # No NS record is found, we need to return with response packet we already have
         if not response_packet.nameserver_records:
@@ -58,18 +54,11 @@ def __perform_recursive_lookup(qpacket: DNSPacket) -> DNSPacket:
             return response_packet
         # Try to find a 'NS' type record with a corresponding 'A' type record in the additional section
         # If found, switch Nameserver and retry the loop i.e perform the lookup on new NameServer again
-        ns_record_set: set[str] = {
-            ns_rec.nsdname for ns_rec in response_packet.nameserver_records
-        }
-        additional_records: List[Union[A, AAAA, NS]] = (
-            response_packet.additional_records
-        )
+        ns_record_set: set[str] = {ns_rec.nsdname for ns_rec in response_packet.nameserver_records}
+        additional_records: List[Union[A, AAAA, NS]] = response_packet.additional_records
         if additional_records:
             for ad_rec in additional_records:
-                if (
-                    ad_rec.rtype.value == RecordType.A.value
-                    and ad_rec.name in ns_record_set
-                ):
+                if ad_rec.rtype.value == RecordType.A.value and ad_rec.name in ns_record_set:
                     server_addr = str(ad_rec.address)
                     break
         else:
@@ -83,9 +72,7 @@ def __perform_recursive_lookup(qpacket: DNSPacket) -> DNSPacket:
                         question_count=1,
                         is_recursion_desired=True,
                     ),
-                    questions=[
-                        Question(ns_record.nsdname, RecordType.A, RecordClass.IN)
-                    ],
+                    questions=[Question(ns_record.nsdname, RecordType.A, RecordClass.IN)],
                 )
             )
             # No 'A' Type record is found, we need to return with response packet we already have
@@ -96,17 +83,11 @@ def __perform_recursive_lookup(qpacket: DNSPacket) -> DNSPacket:
             server_addr = str(a_type_record.address)
 
 
-def handle_request(
-    master_socket: socket.socket, received_bytes: bytes, return_address: tuple[str, int]
-) -> None:
+def handle_request(master_socket: socket.socket, received_bytes: bytes, return_address: tuple[str, int]) -> None:
     query_packet: DNSPacket = DNSParser(bytearray(received_bytes)).get_dns_packet()
     # TODO: Send query for each question in query_packet
-    log(
-        f"Received query for {query_packet.questions[0].name} TYPE {query_packet.questions[0].rtype}"
-    )
+    log(f"Received query for {query_packet.questions[0].name} TYPE {query_packet.questions[0].rtype}")
     response_packet: DNSPacket = __perform_recursive_lookup(query_packet)
     response_packet.header.is_recursion_available = True
     master_socket.sendto(response_packet.to_bin(), return_address)
-    log(
-        f"Query for {query_packet.questions[0].name} TYPE {query_packet.questions[0].rtype} successfully processed"
-    )
+    log(f"Query for {query_packet.questions[0].name} TYPE {query_packet.questions[0].rtype} successfully processed")

optimus/iterator/iter.py

@@ -0,0 +1,19 @@
+class ByteArrayIterator:
+    def __init__(self, data: bytearray) -> None:
+        self.__data = data
+        self.__ptr = 0
+
+    def seek_ptr_pos(self, pos: int) -> None:
+        self.__ptr = pos
+
+    def get_cur_ptr_pos(self) -> int:
+        return self.__ptr
+
+    def get_n_bytes(self, n: int) -> bytearray:
+        cur_ptr_pos = self.get_cur_ptr_pos()
+        return self.__data[cur_ptr_pos : cur_ptr_pos + n]
+
+    def get_n_bytes_and_move(self, n: int) -> bytearray:
+        data = self.get_n_bytes(n)
+        self.seek_ptr_pos(self.get_cur_ptr_pos() + n)
+        return data

optimus/optimus_server/server.py → optimus/optimus_server/udp_listener.py

@@ -5,7 +5,7 @@
 from optimus.logging_config.logger import log
 
 
-def run_server(port: int, worker_threads: int):
+def run_udp_listener(port: int, worker_threads: int):
     master_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     master_socket.bind(("0.0.0.0", port))
     log(f"Started Optimus Server on Port {port}")

pyproject.toml

@@ -11,3 +11,8 @@ exclude = [
     'build',  # TOML literal string (single-quotes, no escaping necessary)
     '^\..*',  # TOML basic string (double-quotes, backslash and other characters need escaping)
 ]
+
+[tool.black]
+line-length = 120
+target-version = ['py39']
+