first pass

1 files changed, 47 insertions, 23 deletions

diff --git a/asn.py b/asn.py
index eb5d0a0..eca0524 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,15 @@ 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)
+        if v4:
+            self._listen_v4(v4, port)
+        if v6:
+            self._listen_v6(v6, port)
 
-    def _listen(self, host, port):
+    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 +135,17 @@ class WHOISListener(Common):
                 threading.Thread(target=self._handler,
                         args=(conn,addr,), daemon=True).start()
 
+    def _listen_v6(self, host, port):
+        with (socket.socket(socket.AF_INET6, socket.SOCK_DGRAM,
+                            socket.IPPROTO_UDP) 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:
@@ -305,18 +327,18 @@ if __name__ == '__main__':
     desc = 'asn: map hosts to their corresponding ASN via 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 +364,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()