Refactor Codebase (#10)

* Refactor and improve code quality

* More changes

* Remove commented code

* Rearrange

* Fix issues

optimus/cli/optimus.py

@@ -2,7 +2,7 @@
 from argparse import ArgumentParser
 
 from optimus.__version__ import VERSION
-from optimus.optimus_server.server import run_server
+from optimus.optimus_server.udp import run_udp_listener
 
 
 def main(argv):
@@ -32,7 +32,7 @@ def main(argv):
     arg_parser.add_argument("-v", action="store_true", help="Get version info")
     args = arg_parser.parse_args(argv)
     if args.r:
-        run_server(args.p, args.t)
+        run_udp_listener(args.p, args.t)
     elif args.v:
         print(f"Optimus Version: {VERSION}")
     else:

optimus/dns/dns_packet.py → optimus/dns/models/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.models.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 → optimus/dns/models/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/parser/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