More changes

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/handlers.py

@@ -3,11 +3,11 @@
 import socket
 from typing import List, Union
 
-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.logging_config.logger import log
-from optimus.networking.udp_utils import query_server_over_udp
+from optimus.dns.models.packet import DNSHeader, DNSPacket, Question, ResponseCode
+from optimus.dns.parser.parse import DNSParser
+from optimus.dns.models.records import AAAA, NS, A, Record, RecordClass, RecordType
+from optimus.logging.logger import log
+from optimus.networking.udp import query_server_over_udp
 
 # TODO: Move to a config file and read from file instead of hardcoding
 ROOT_SERVERS = [

optimus/dns/dns_packet.py → optimus/dns/models/packet.py

@@ -1,7 +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

optimus/iterator/iter.py → optimus/dns/parser/iter.py

@@ -1,4 +1,4 @@
-class ByteArrayIterator:
+class BytearrayIterator:
     def __init__(self, data: bytearray) -> None:
         self.__data = data
         self.__ptr = 0

optimus/dns/dns_parser.py → optimus/dns/parser/parse.py

@@ -1,9 +1,9 @@
 from ipaddress import IPv4Address, IPv6Address
 from typing import List, Union
 
-from optimus.bin_data_reader.bin_reader import ByteArrayIterator
-from optimus.dns.dns_packet import DNSHeader, DNSPacket, Question, ResponseCode
-from optimus.dns.dns_records import (
+from optimus.dns.parser.iter import BytearrayIterator
+from optimus.dns.models.packet import DNSHeader, DNSPacket, Question, ResponseCode
+from optimus.dns.models.records import (
     AAAA,
     CNAME,
     MX,
@@ -21,7 +21,7 @@ class DNSParser:
     def __init__(self, bin_data: bytearray) -> None:
         if not bin_data:
             raise Exception("No binary data given to parse")
-        self.__iter = ByteArrayIterator(bin_data)
+        self.__iter = BytearrayIterator(bin_data)
 
     def __to_int(self, data: bytearray) -> int:
         value = 0
@@ -48,7 +48,7 @@ def get_dns_packet(self) -> DNSPacket:
         return DNSPacket(dns_header, questions, answers, nameserver_records, additional_records)
 
     def __parse_record_name(self) -> str:
-        msb_data: int = self.__iter.parse_bytes_to_int(1)
+        msb_data: int = self.__to_int(self.__iter.get_n_bytes(1))
         if self.__is_label_compressed(msb_data):
             # In this case, the next byte should be treated as OFFSET
             # which specifies the position from where the name has to be
@@ -64,7 +64,7 @@ def __is_label_compressed(self, label_length) -> bool:
         return bool(label_length ^ 0xC0 == 0)
 
     def __parse_compressed_label(self) -> str:
-        offset: int = self.__iter.get_n_bytes_and_move(1)
+        offset: int = self.__to_int(self.__iter.get_n_bytes_and_move(1))
         cur_ptr_pos: int = self.__iter.get_cur_ptr_pos()
         self.__iter.seek_ptr_pos(offset)
         name: str = self.__parse_sequence_of_labels()
@@ -78,19 +78,19 @@ def __parse_sequence_of_labels(self) -> str:
         # 2 MSB bits of the this label_length field are always 0
         # So the label_length can have values between 0 and 63, meaning
         # that the name can take only between 0-63 octets
-        label_length: int = self.__iter.get_n_bytes_and_move(1)
+        label_length: int = self.__to_int(self.__iter.get_n_bytes_and_move(1))
         # Parse sequence of labels to decode the Name
         while label_length != 0:
             if self.__is_label_compressed(label_length):
                 name_str = self.__parse_compressed_label()
                 full_name.append(name_str)
                 return ".".join(full_name)
             for _ in range(0, label_length):
-                label: int = self.__iter.get_n_bytes_and_move(1)
+                label: int = self.__to_int(self.__iter.get_n_bytes_and_move(1))
                 name_str = name_str + chr(label)
             full_name.append(name_str)
             name_str = ""
-            label_length = self.__iter.get_n_bytes_and_move(1)
+            label_length = self.__to_int(self.__iter.get_n_bytes_and_move(1))
         return ".".join(full_name)
 
     def __get_dns_header(self) -> DNSHeader:
@@ -139,7 +139,7 @@ def __get_dns_header(self) -> DNSHeader:
 
     def __get_ques_section(self, total_questions) -> List[Question]:
         questions: List[Question] = []
-        for _ in range(0, total_questions):
+        for _ in range(total_questions):
             # Parse name
             name: str = self.__parse_record_name()
             # Parse type
@@ -168,13 +168,13 @@ def __read_response_records(
             # rclass = RecordClass.from_value(self.__iter.get_n_bytes_and_move(2))
             # ttl = self.__iter.get_n_bytes_and_move(4)
             # length = self.__iter.get_n_bytes_and_move(2)
-            ipv4_addr_int: int = self.__iter.get_n_bytes_and_move(4)
+            ipv4_addr_int: int = self.__to_int(self.__iter.get_n_bytes_and_move(4))
             record = A(name, rtype, rclass, ttl, length, IPv4Address(ipv4_addr_int))
         elif rtype.value == RecordType.AAAA.value:
             # rclass = RecordClass.from_value(self.__iter.get_n_bytes_and_move(2))
             # ttl = self.__iter.get_n_bytes_and_move(4)
             # length = self.__iter.get_n_bytes_and_move(2)
-            ipv6_addr_int: int = self.__iter.get_n_bytes_and_move(16)
+            ipv6_addr_int: int = self.__to_int(self.__iter.get_n_bytes_and_move(16))
             record = AAAA(name, rtype, rclass, ttl, length, IPv6Address(ipv6_addr_int))
         elif rtype.value == RecordType.CNAME.value:
             # rclass = RecordClass.from_value(self.__iter.get_n_bytes_and_move(2))

optimus/networking/udp.py

@@ -1,6 +1,6 @@
 import socket
 
-from optimus.logging_config.logger import log_error
+from optimus.logging.logger import log_error
 
 
 def query_server_over_udp(bin_data: bytearray, server_addr: str) -> bytes:

optimus/optimus_server/udp_listener.py → optimus/optimus_server/udp.py

@@ -2,7 +2,7 @@
 from concurrent import futures
 
 from optimus.dns.handlers import handle_request
-from optimus.logging_config.logger import log
+from optimus.logging.logger import log
 
 
 def run_udp_listener(port: int, worker_threads: int):

setup.py

@@ -4,13 +4,13 @@
 
 here = os.path.abspath(os.path.dirname(__file__))
 
-META: dict = {}
+META: dict[str, str] = {}
 with open(os.path.join(here, "optimus", "__version__.py")) as f:
     exec(f.read(), META)
 
 setup(
     name="optimus",
-    version=META.get("VERSION"),
+    version=META["VERSION"],
     description="A toy DNS server",
     author="Anchal",
     author_email="[email protected]",