From 3c872618ecb6d78b4001b9c47baaace568890603 Mon Sep 17 00:00:00 2001
From: Martin Braun <braun@neuroforge.de>
Date: Sun, 13 Oct 2024 21:56:53 +0200
Subject: [PATCH] improve code in dns.py, improve stack.yml deployment to
 something that works

---
 deploy/stack.yml | 12 +++++++++++
 dns.py           | 53 +++++++++++++++++++++++++++++-------------------
 2 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/deploy/stack.yml b/deploy/stack.yml
index 791a359..8da7b32 100644
--- a/deploy/stack.yml
+++ b/deploy/stack.yml
@@ -9,6 +9,9 @@ services:
       S3_REFRESH_INTERVAL: "10"
       DOCKER_NETWORK_INFO_CACHE_REFRESH_INTERVAL: "10"
       STRIP_DOMAIN_ENDINGS: ".localdomain.,.docker.,.docker.localdomain."
+    dns:
+      - "8.8.8.8"
+      - "8.8.4.4"
     networks:
       - host
     volumes:
@@ -25,6 +28,9 @@ services:
     environment:
       SWARM_NODE_ID: "{{ .Node.ID }}"
       EXPORTER_INTERVAL: "10"
+    dns:
+      - "8.8.8.8"
+      - "8.8.4.4"
     volumes:
       - /var/run/docker.sock:/var/run/docker.sock
     deploy:
@@ -39,6 +45,9 @@ services:
       - .env
     environment:
       NODES_INTERVAL: "60"
+    dns:
+      - "8.8.8.8"
+      - "8.8.4.4"
     volumes:
       - /var/run/docker.sock:/var/run/docker.sock
     deploy:
@@ -54,6 +63,9 @@ services:
     command: ["python", "-u", "merger.py"]
     env_file:
       - .env
+    dns:
+      - "8.8.8.8"
+      - "8.8.4.4"
     environment:
       MERGE_INTERVAL: "10"
 
diff --git a/dns.py b/dns.py
index 4938e0a..fb98ce6 100644
--- a/dns.py
+++ b/dns.py
@@ -297,6 +297,8 @@ def resolve_dnsA_to_ip(network_data, networks, domain):
 
     print_debug(f"Resolved DNS A records: {dnsA_records}")
     return list(dnsA_records)
+
+
     
 # DNS Server
 class DNSServer:
@@ -311,6 +313,34 @@ def __init__(self, ip="0.0.0.0", port=53, upstream_dns="8.8.8.8"):
         self.upstream_dns = upstream_dns
         print_debug(f"DNS server initialized on {self.ip}:{self.port}")
 
+    def forward_dns_request(self, request):
+        """
+        Forwards a DNS request to an upstream DNS server and returns the response.
+        """
+        # fallback to upstream DNS server
+        try:
+            # Convert the request to binary format
+            query_data = request.pack()
+
+            # Create a socket and send the query to the upstream server
+            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            sock.settimeout(2)
+
+            # Send the request to the upstream DNS server
+            sock.sendto(query_data, (self.upstream_dns, 53))
+            
+            # Receive the response from the upstream server
+            data, _ = sock.recvfrom(4096)
+            sock.close()
+
+            # Parse the response
+            reply = DNSRecord.parse(data)
+        except Exception as e:
+            reply = request.reply()
+            reply.header.rcode = RCODE.SERVFAIL
+
+        return reply
+
     def handle_request(self, data, addr):
         # Parse incoming DNS request
         request = DNSRecord.parse(data)
@@ -344,27 +374,8 @@ def handle_request(self, data, addr):
             dnsA_records = resolve_dnsA_to_ip(network_data, networks, domain)
 
             if len(dnsA_records) == 0:
-                # fallback to upstream DNS server
-                try:
-                    # Convert the request to binary format
-                    query_data = request.pack()
-
-                    # Create a socket and send the query to the upstream server
-                    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-                    sock.settimeout(2)
-
-                    # Send the request to the upstream DNS server
-                    sock.sendto(query_data, (self.upstream_dns, 53))
-                    
-                    # Receive the response from the upstream server
-                    data, _ = sock.recvfrom(4096)
-                    sock.close()
-
-                    # Parse the response
-                    reply = DNSRecord.parse(data)
-                except Exception as e:
-                    reply = request.reply()
-                    reply.header.rcode = RCODE.SERVFAIL
+                print_debug(f"No DNS A records found for domain {domain}. Falling back to upstream DNS server.")
+                reply = self.forward_dns_request(request)
             else:
                 reply = DNSRecord(DNSHeader(id=request.header.id, qr=1, aa=1, ra=0), q=request.q)
                 for ip in dnsA_records: