Cleanup, implement zero padding IP octets, switch to fstrings

sudomemoDNS.py

@@ -12,20 +12,30 @@ import requests
 import json
 import sys
 
+
 def get_platform():
     platforms = {
-        'linux1' : 'Linux',
-        'linux2' : 'Linux',
-        'darwin' : 'macOS',
-        'win32' : 'Windows'
+        'linux1': 'Linux',
+        'linux2': 'Linux',
+        'darwin': 'macOS',
+        'win32': 'Windows'
     }
     if sys.platform not in platforms:
         return sys.platform
 
     return platforms[sys.platform]
 
+
 SUDOMEMODNS_VERSION = "1.2"
 
+
+# Adds preceding zeros to IP addresses
+# e.g `1.1.1.1` becomes `001.001.001.001`
+def format_ip(address):
+    octets = str(address).split(".")
+    return f"{int(octets[0]):02d}.{int(octets[1]):02d}.{int(octets[2]):02d}.{int(octets[3]):02d}"
+
+
 def get_ip():
     s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     try:
@@ -38,25 +48,27 @@ def get_ip():
         s.close()
     return IP
 
+
 EPOCH = datetime(1970, 1, 1)
 SERIAL = int((datetime.utcnow() - EPOCH).total_seconds())
 MY_IP = get_ip()
 
 print("+===============================+")
 print("|      Sudomemo DNS Server      |")
-print("|          Version " + SUDOMEMODNS_VERSION + "          |")
+print(f"|          Version {SUDOMEMODNS_VERSION}          |")
 print("+===============================+\n")
 
 print("== Welcome to sudomemoDNS! ==")
-print("This server will allow you to connect to Sudomemo when your Internet Service Provider does not work with custom DNS.\n")
+print(
+    "This server will allow you to connect to Sudomemo when your Internet Service Provider does not work with custom DNS.\n")
 
 print("== How To Use ==")
 print("First, make sure that your console is connected to the same network as this computer.\n")
 
 print("Then, put these settings in for DNS on your console:")
 print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
-print("Primary DNS:  ",MY_IP)
-print("Secondary DNS: 8.8.8.8")
+print(f"Primary DNS:  {MY_IP}")
+print("Secondary DNS: 008.008.008.008")
 print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n")
 
 print("== Getting Help ==")
@@ -74,24 +86,32 @@ TYPE_LOOKUP = {
     TXT: QTYPE.TXT,
 }
 
+
 # Can't seem to turn off DNSLogger with a None type so let's just null it out with a dummy function
 
 class SudomemoDNSLogger(object):
     def log_recv(self, handler, data):
         pass
+
     def log_send(self, handler, data):
         pass
+
     def log_request(self, handler, request):
         print("[INFO] Received DNS request from console at " + handler.client_address[0])
+
     def log_reply(self, handler, reply):
         print("[INFO] Sent response to console at " + handler.client_address[0])
+
     def log_error(self, handler, e):
         print("[ERROR] Invalid DNS request from " + handler.client_address[0])
+
     def log_truncated(self, handler, reply):
         pass
+
     def log_data(self, dnsobj):
         pass
 
+
 class Record:
     def __init__(self, rdata_type, *args, rtype=None, rname=None, ttl=None, **kwargs):
         if isinstance(rdata_type, RD):
@@ -144,12 +164,13 @@ class Record:
 ZONES = {}
 
 try:
-  get_zones = requests.get("https://www.sudomemo.net/api/dns_zones.json", headers={'User-Agent': 'SudomemoDNS/' + SUDOMEMODNS_VERSION + ' (' + get_platform() + ')'})
+    get_zones = requests.get("https://www.sudomemo.net/api/dns_zones.json",
+                             headers={'User-Agent': 'SudomemoDNS/' + SUDOMEMODNS_VERSION + ' (' + get_platform() + ')'})
 except requests.exceptions.Timeout:
     print("[ERROR] Unable to load DNS data: Connection to Sudomemo timed out. Are you connected to the Internet?")
 except requests.exceptions.RequestException as e:
     print("[ERROR] Unable load DNS data.")
-  print("[ERROR] Exception: ",e)
+    print("[ERROR] Exception: ", e)
     sys.exit(1)
 try:
     zones = json.loads(get_zones.text)
@@ -158,12 +179,13 @@ except ValueError as e:
 
 for zone in zones:
     if zone["type"] == "a":
-    ZONES[zone["name"]] = [ Record(A, zone["value"]) ]
+        ZONES[zone["name"]] = [Record(A, zone["value"])]
     elif zone["type"] == "p":
-    ZONES[zone["name"]] = [ Record(A, socket.gethostbyname(zone["value"])) ]
+        ZONES[zone["name"]] = [Record(A, socket.gethostbyname(zone["value"]))]
 
 print("[INFO] DNS information loaded successfully.")
 
+
 class Resolver:
     def __init__(self):
         self.zones = {DNSLabel(k): v for k, v in ZONES.items()}
@@ -171,6 +193,7 @@ class Resolver:
     def resolve(self, request, handler):
         reply = request.reply()
         zone = self.zones.get(request.q.qname)
+
         if zone is not None:
             for zone_records in zone:
                 rr = zone_records.try_rr(request.q)
@@ -210,13 +233,13 @@ elif get_platform() == 'Windows':
 try:
     servers = [
         DNSServer(resolver=resolver, port=53, address=MY_IP, tcp=True, logger=dnsLogger),
-    DNSServer(resolver=resolver, port=53, address=MY_IP, tcp=False, logger=dnsLogger),
+        DNSServer(resolver=resolver, port=53, address=MY_IP, tcp=False, logger=dnsLogger)
     ]
+
 except PermissionError:
     print("[ERROR] Permission error: Check that you are running this as an administrator or root")
     sys.exit(1)
 
-
 print("[INFO] sudomemoDNS is ready. Now waiting for DNS requests from your console...")
 
 if __name__ == '__main__':