From 1913ba59b1f6ad4677f9dc7bb41e5774225eaaee Mon Sep 17 00:00:00 2001
From: Michal Charemza <michal@charemza.name>
Date: Thu, 16 Jan 2020 15:35:48 +0000
Subject: [PATCH] (feat) Add basic rewriting

---
 README.md          |  5 ++++-
 dnsrewriteproxy.py | 12 +++++++++++-
 test.py            | 41 +++++++++++++++++++++++++++++------------
 3 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 59b84fe..d25ffbe 100644
--- a/README.md
+++ b/README.md
@@ -5,10 +5,13 @@ A DNS proxy server that conditionally rewrites and filters A record requests
 
 ## Usage
 
+By default the proxy will listen on port 53, and proxy requests to the servers in `/etc/resolve.conf`. However, by default all requests are blocked without explicit rules, so to proxy requests you must configure at least one rewrite rule.
+
 ```python
 from dnsrewriteproxy import DnsProxy
 
-start = DnsProxy()
+# Proxy all incoming A record requests without any rewriting
+start = DnsProxy(rules=((r'(^.*$)', r'\1'),))
 
 # Proxy is running, accepting UDP requests on port 53
 stop = await start()
diff --git a/dnsrewriteproxy.py b/dnsrewriteproxy.py
index 5358f53..e10506e 100644
--- a/dnsrewriteproxy.py
+++ b/dnsrewriteproxy.py
@@ -19,6 +19,7 @@ from enum import (
     IntEnum,
 )
 import logging
+import re
 import socket
 
 from aiodnsresolver import (
@@ -61,6 +62,7 @@ def DnsProxy(
         get_resolver=get_resolver_default, get_logger=get_logger_default,
         get_socket=get_socket_default,
         num_workers=1000, downstream_queue_maxsize=10000, upstream_queue_maxsize=10000,
+        rules=(),
 ):
 
     class ERRORS(IntEnum):
@@ -165,8 +167,16 @@ def DnsProxy(
         name_bytes = query.qd[0].name
         name_str = query.qd[0].name.decode('idna')
 
+        for pattern, replace in rules:
+            rewritten_name_str, num_matches = re.subn(pattern, replace, name_str)
+            if num_matches:
+                break
+        else:
+            # No break was triggered, i.e. no match
+            return error(query, ERRORS.REFUSED)
+
         try:
-            ip_addresses = await resolve(name_str, TYPES.A)
+            ip_addresses = await resolve(rewritten_name_str, TYPES.A)
         except DnsRecordDoesNotExist:
             return error(query, ERRORS.NXDOMAIN)
         except DnsResponseCode as dns_response_code_error:
diff --git a/test.py b/test.py
index a84dfcb..a4e8531 100644
--- a/test.py
+++ b/test.py
@@ -7,6 +7,7 @@ from aiodnsresolver import (
     TYPES,
     Resolver,
     IPv4AddressExpiresAt,
+    DnsResponseCode,
 )
 from dnsrewriteproxy import (
     DnsProxy,
@@ -26,23 +27,39 @@ class TestProxy(unittest.TestCase):
         self.addCleanup(asyncio.get_running_loop().run_until_complete, coroutine())
 
     @async_test
-    async def test_e2e(self):
-
-        def get_socket():
-            sock = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
-            sock.setblocking(False)
-            sock.bind(('', 3535))
-            return sock
+    async def test_e2e_no_match_rule(self):
+        resolve, clear_cache = get_resolver()
+        self.add_async_cleanup(clear_cache)
+        start = DnsProxy(get_socket=get_socket)
+        stop = await start()
+        self.add_async_cleanup(stop)
 
-        async def get_nameservers(_, __):
-            for _ in range(0, 5):
-                yield (0.5, ('127.0.0.1', 3535))
+        with self.assertRaises(DnsResponseCode):
+            await resolve('www.google.com', TYPES.A)
 
-        resolve, clear_cache = Resolver(get_nameservers=get_nameservers)
+    @async_test
+    async def test_e2e_match_all(self):
+        resolve, clear_cache = get_resolver()
         self.add_async_cleanup(clear_cache)
-        start = DnsProxy(get_socket=get_socket)
+        start = DnsProxy(get_socket=get_socket, rules=((r'(^.*$)', r'\1'),))
         stop = await start()
         self.add_async_cleanup(stop)
 
         response = await resolve('www.google.com', TYPES.A)
+
         self.assertTrue(isinstance(response[0], IPv4AddressExpiresAt))
+
+
+def get_socket():
+    sock = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
+    sock.setblocking(False)
+    sock.bind(('', 3535))
+    return sock
+
+
+def get_resolver():
+    async def get_nameservers(_, __):
+        for _ in range(0, 5):
+            yield (0.5, ('127.0.0.1', 3535))
+
+    return Resolver(get_nameservers=get_nameservers)