Almost-dynamic zone update

dnserver/main.py

@@ -10,7 +10,7 @@
 from dnslib.proxy import ProxyResolver as LibProxyResolver
 from dnslib.server import BaseResolver as LibBaseResolver, DNSServer as LibDNSServer
 
-from .load_records import Zone, load_records
+from .load_records import Records, Zone, load_records
 
 __all__ = 'DNSServer', 'logger'
 
@@ -107,10 +107,8 @@ def resolve(request, handler, records):
 
 
 class BaseResolver(LibBaseResolver):
-    def __init__(self, zones_file: str | Path):
-        records = load_records(zones_file)
+    def __init__(self, records: Records):
         self.records = [Record(zone) for zone in records.zones]
-        logger.info('loaded %d zone record from %s, without proxy DNS server', len(self.records), zones_file)
         super().__init__()
 
     def resolve(self, request, handler):
@@ -124,12 +122,8 @@ def resolve(self, request, handler):
 
 
 class ProxyResolver(LibProxyResolver):
-    def __init__(self, zones_file: str | Path, upstream: str):
-        records = load_records(zones_file)
+    def __init__(self, records: Records, upstream: str):
         self.records = [Record(zone) for zone in records.zones]
-        logger.info(
-            'loaded %d zone record from %s, with %s as a proxy DNS server', len(self.records), zones_file, upstream
-        )
         super().__init__(address=upstream, port=53, timeout=5)
 
     def resolve(self, request, handler):
@@ -146,18 +140,26 @@ class DNSServer:
     def __init__(
         self, zones_file: str | Path, *, port: int | str | None = DEFAULT_PORT, upstream: str | None = DEFAULT_UPSTREAM
     ):
-        self.zones_file = zones_file
         self.port: int = DEFAULT_PORT if port is None else int(port)
         self.upstream: str | None = upstream
         self.udp_server: LibDNSServer | None = None
         self.tcp_server: LibDNSServer | None = None
 
+        self.records: Records = load_records(zones_file)
+        logger.info(
+            'loaded %d zone record from %s, with %s as a proxy DNS server',
+            len(self.records.zones),
+            zones_file,
+            upstream,
+        )
+
     def start(self):
-        logger.info('starting DNS server on port %d, upstream DNS server "%s"', self.port, self.upstream)
         if self.upstream:
-            resolver = ProxyResolver(self.zones_file, self.upstream)
+            logger.info('starting DNS server on port %d, upstream DNS server "%s"', self.port, self.upstream)
+            resolver = ProxyResolver(self.records, self.upstream)
         else:
-            resolver = BaseResolver(self.zones_file)
+            logger.info('starting DNS server on port %d, without upstream DNS server', self.port)
+            resolver = BaseResolver(self.records)
 
         self.udp_server = LibDNSServer(resolver, port=self.port)
         self.tcp_server = LibDNSServer(resolver, port=self.port, tcp=True)
@@ -173,3 +175,8 @@ def stop(self):
     @property
     def is_running(self):
         return (self.udp_server and self.udp_server.isAlive()) or (self.tcp_server and self.tcp_server.isAlive())
+
+    def add_record(self, **kwargs):
+        zone = Zone(**kwargs)
+        self.records.zones.append(zone)
+        return zone

tests/test_dnsserver.py

@@ -1,5 +1,6 @@
 from typing import Any, Callable, Dict, List
 
+import dns
 import pytest
 from dirty_equals import IsIP, IsPositive
 from dns.resolver import NoAnswer, Resolver as RawResolver
@@ -33,26 +34,28 @@ def convert_answer(answer) -> Dict[str, Any]:
 
 
 @pytest.fixture(scope='session')
-def dns_resolver():
+def dns_server():
     port = 5053
 
     server = DNSServer('example_zones.toml', port=port)
     server.start()
-
     assert server.is_running
+    yield server
+    server.stop
+
 
+@pytest.fixture(scope='session')
+def dns_resolver(dns_server: DNSServer):
     resolver = RawResolver()
     resolver.nameservers = ['127.0.0.1']
-    resolver.port = port
+    resolver.port = dns_server.port
 
     def resolve(name: str, type_: str) -> List[Dict[str, Any]]:
         answers = resolver.resolve(name, type_)
         return [convert_answer(answer) for answer in answers]
 
     yield resolve
 
-    server.stop()
-
 
 def test_a_record(dns_resolver: Resolver):
     assert dns_resolver('example.com', 'A') == [
@@ -155,3 +158,31 @@ def resolve(name: str, type_: str) -> List[Dict[str, Any]]:
             resolve('python.org', 'A')
     finally:
         server.stop()
+
+
+def test_dynamic_zone_update(dns_server: DNSServer, dns_resolver: Resolver):
+    assert dns_resolver('example.com', 'A') == [
+        {
+            'type': 'A',
+            'value': '1.2.3.4',
+        },
+    ]
+    with pytest.raises(dns.resolver.NXDOMAIN):
+        dns_resolver('another-example.org', 'A')
+
+    dns_server.stop()
+    dns_server.add_record(host='another-example.com', type='A', answer='2.3.4.5')
+    dns_server.start()
+
+    assert dns_resolver('example.com', 'A') == [
+        {
+            'type': 'A',
+            'value': '1.2.3.4',
+        },
+    ]
+    assert dns_resolver('another-example.com', 'A') == [
+        {
+            'type': 'A',
+            'value': '2.3.4.5',
+        },
+    ]