asn: support dualstack HTTP/WHOIS sockets

2 files changed, 64 insertions, 35 deletions

diff --git a/README b/README
index cd58d79..c85782e 100644
--- a/README
+++ b/README
@@ -1,18 +1,22 @@
 asn: map hosts to their corresponding ASN via HTTP/WHOIS
 
-usage: asn.py [-h] [--whois-host WHOIS_HOST] [--whois-port WHOIS_PORT] [--http-host HTTP_HOST]
-              [--http-port HTTP_PORT] [--update] [--populate]
+usage: asn.py [-h] [--bind-whois-v4 WHOIS_V4] [--bind-whois-v6 WHOIS_V6] [--bind-whois-port WHOIS_PORT]
+              [--bind-http-v4 HTTP_V4] [--bind-http-v6 HTTP_V6] [--bind-http-port HTTP_PORT] [--update] [--populate]
 
 options:
   -h, --help            show this help message and exit
-  --whois-host WHOIS_HOST
-                        IP to listen on for WHOIS service (default: 127.0.0.1)
-  --whois-port WHOIS_PORT
-                        Port to listen on for WHOIS service (default: 4343)
-  --http-host HTTP_HOST
-                        IP to listen on for HTTP service (default: 127.0.0.1)
-  --http-port HTTP_PORT
-                        Port to listen on for HTTP service (default: 8080)
+  --bind-whois-v4 WHOIS_V4
+                        WHOIS IPv4 host (default: 127.0.0.1)
+  --bind-whois-v6 WHOIS_V6
+                        WHOIS IPv6 host (default: ::1)
+  --bind-whois-port WHOIS_PORT
+                        WHOIS port (default: 4343)
+  --bind-http-v4 HTTP_V4
+                        HTTP IPv4 host (default: 127.0.0.1)
+  --bind-http-v6 HTTP_V6
+                        HTTP IPv6 host (default: ::1)
+  --bind-http-port HTTP_PORT
+                        HTTP port (default: 8080)
   --update              Update dataset submodule and create/populate cache (default: False)
   --populate            Create and populate cache from current dataset (default: False)
 
diff --git a/asn.py b/asn.py
index eb5d0a0..cec4017 100755
--- a/asn.py
+++ b/asn.py
@@ -60,13 +60,21 @@ class Common:
         return announcements
 
 class HTTPListener(Common):
-    def __init__(self, db, host, port):
+    def __init__(self, db, v4, v6, port):
         Common.__init__(self)
         self.db = db
 
         self._app = FlaskWrapper("asn")
         self._app.add_endpoint(endpoint="/", name="index", handler=self._handler)
-        self._app.run(host=host, port=port)
+
+        if v4 and v6:
+            listen = f"{v4}:{port} {v6}:{port}"
+        elif v4:
+            listen = f"{v4}:{port}"
+        elif v6:
+            listen = f"{v6}:{port}"
+
+        self._app.run(listen=listen)
 
     def _handler(self):
         if request.headers.getlist("X-Forwarded-For"):
@@ -109,12 +117,17 @@ class HTTPHandler:
         return res
 
 class WHOISListener(Common):
-    def __init__(self, db, host, port):
+    def __init__(self, db, v4, v6, port):
         Common.__init__(self)
         self.db = db
-        self._listen(host, port)
-
-    def _listen(self, host, port):
+        if v4:
+            listen_v4 = Thread(target=self._listen_v4, args=(v4, port))
+            listen_v4.start()
+        if v6:
+            listen_v6 = Thread(target=self._listen_v6, args=(v6, port))
+            listen_v6.start()
+
+    def _listen_v4(self, host, port):
         with socket.socket() as _socket:
             _socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
             _socket.bind((host, port))
@@ -124,6 +137,16 @@ class WHOISListener(Common):
                 threading.Thread(target=self._handler,
                         args=(conn,addr,), daemon=True).start()
 
+    def _listen_v6(self, host, port):
+        with socket.socket(family=socket.AF_INET6) as _socket:
+            _socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            _socket.bind((host, port))
+            _socket.listen()
+            while True:
+                conn, addr = _socket.accept()
+                threading.Thread(target=self._handler,
+                        args=(conn,addr,), daemon=True).start()
+
     def _handler(self, conn, addr):
         resp = ''
         try:
@@ -302,21 +325,21 @@ class DB:
                 break
 
 if __name__ == '__main__':
-    desc = 'asn: map hosts to their corresponding ASN via WHOIS'
+    desc = 'asn: map hosts to their corresponding ASN via HTTP/WHOIS'
     parser = argparse.ArgumentParser(description=desc,
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    parser.add_argument('--whois-host', dest='whois_host', type=str, action='store',
-                        help='IP to listen on for WHOIS service',
-                        default="127.0.0.1", required=False)
-    parser.add_argument('--whois-port', dest='whois_port', type=int, action='store',
-                        help='Port to listen on for WHOIS service',
-                        default=4343, required=False)
-    parser.add_argument('--http-host', dest='http_host', type=str, action='store',
-                        help='IP to listen on for HTTP service',
-                        default="127.0.0.1", required=False)
-    parser.add_argument('--http-port', dest='http_port', type=int, action='store',
-                        help='Port to listen on for HTTP service', default=8080,
-                        required=False)
+    parser.add_argument('--bind-whois-v4', dest='whois_v4', type=str, action='store',
+                        default="127.0.0.1", required=False, help='WHOIS IPv4 host')
+    parser.add_argument('--bind-whois-v6', dest='whois_v6', type=str, action='store',
+                        default="::1", required=False, help='WHOIS IPv6 host')
+    parser.add_argument('--bind-whois-port', dest='whois_port', type=int, action='store',
+                        default=4343, required=False, help='WHOIS port')
+    parser.add_argument('--bind-http-v4', dest='http_v4', type=str, action='store',
+                        default="127.0.0.1", required=False, help='HTTP IPv4 host')
+    parser.add_argument('--bind-http-v6', dest='http_v6', type=str, action='store',
+                        default="::1", required=False, help='HTTP IPv6 host')
+    parser.add_argument('--bind-http-port', dest='http_port', type=int, action='store',
+                        default=8080, required=False, help='HTTP port')
     parser.add_argument('--update', dest='update', action='store_true',
                         help='Update dataset submodule and create/populate cache',
                         required=False)
@@ -342,14 +365,16 @@ if __name__ == '__main__':
         else:
             log.info('no changes since last update')
 
-    if args.whois_host and args.whois_port:
-        log.info(f'WHOIS: listening on {args.whois_host}:{args.whois_port}')
-        whois = Thread(target=WHOISListener, args=(db, args.whois_host,
+    if (args.whois_v4 or args.whois_v6) and args.whois_port:
+        log.info(f'WHOIS: listening on {args.whois_port}')
+        whois = Thread(target=WHOISListener, args=(db, args.whois_v4,
+                                                   args.whois_v6,
                                                    args.whois_port))
         whois.start()
 
-    if args.http_host and args.http_port:
-        log.info(f'HTTP: listening on {args.http_host}:{args.http_port}')
-        http = Thread(target=HTTPListener, args=(db, args.http_host,
+    if (args.http_v4 or args.http_v6) and args.http_port:
+        log.info(f'HTTP: listening on {args.http_port}')
+        http = Thread(target=HTTPListener, args=(db, args.http_v4,
+                                                 args.http_v6,
                                                  args.http_port))
         http.start()