From 6ed92c4c5f85bc6975497af868210af072a18168 Mon Sep 17 00:00:00 2001
From: Aaron LI <aaronly.me@outlook.com>
Date: Tue, 8 Nov 2016 19:14:34 +0800
Subject: webui: Add support of controlling the allowed hosts

---
 fg21sim/webui/utils.py     | 35 +++++++++++++++++++++++++++++------
 fg21sim/webui/websocket.py | 25 ++++++++++++++++++-------
 2 files changed, 47 insertions(+), 13 deletions(-)

(limited to 'fg21sim/webui')

diff --git a/fg21sim/webui/utils.py b/fg21sim/webui/utils.py
index fbb2ade..b3b013a 100644
--- a/fg21sim/webui/utils.py
+++ b/fg21sim/webui/utils.py
@@ -3,17 +3,22 @@
 
 """
 Utilities for the Web UI
+------------------------
 
-TODO:
-* Add a function to determine whether the two IPs are in the same sub-network.
-  References:
-  + Stackoverflow: How can I check if an IP is in a network in Python?
-    https://stackoverflow.com/q/819355/4856091
+get_host_ip :
+    Get the IP address of the host extracted from the input URL.
+
+get_local_ip :
+    Get the local IP address of this machine.
+
+ip_in_network :
+    Whether the IP address is contained in the network?
 """
 
 
-from urllib.parse import urlparse
+import ipaddress
 import socket
+from urllib.parse import urlparse
 
 
 def get_host_ip(url):
@@ -76,3 +81,21 @@ def get_local_ip(host="localhost", timeout=3.0):
         except (socket.gaierror, socket.timeout):
             ip = None
     return ip
+
+
+def ip_in_network(ip, network):
+    """
+    Check whether the IP address is contained in the network?
+
+    Parameters
+    ----------
+    ip : `~ipaddress.IPv4Address`, str
+        An `~ipaddress.IPv4Address` instance or a string of the IPv4 address
+    network : `~ipaddress.IPv4Network`, str
+        An `~ipaddress.IPv4Network` instance or a string of the IPv4 network
+    """
+    if not isinstance(ip, ipaddress.IPv4Address):
+        ip = ipaddress.IPv4Address(ip)
+    if not isinstance(network, ipaddress.IPv4Network):
+        network = ipaddress.IPv4Network(network)
+    return ip in network
diff --git a/fg21sim/webui/websocket.py b/fg21sim/webui/websocket.py
index 8c2f427..fffc19c 100644
--- a/fg21sim/webui/websocket.py
+++ b/fg21sim/webui/websocket.py
@@ -22,10 +22,11 @@ import json
 import logging
 
 import tornado.websocket
+from tornado.options import options
 
 from ..configs import ConfigManager
 from ..errors import ConfigError
-from .utils import get_host_ip
+from .utils import get_host_ip, ip_in_network
 
 
 logger = logging.getLogger(__name__)
@@ -73,22 +74,32 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
         Currently, only allow access from the ``localhost``
         (i.e., 127.0.0.1) and local LAN.
         """
+        self.from_localhost = False
         logger.info("WebSocket: {0}: origin: {1}".format(self.name, origin))
         ip = get_host_ip(url=origin)
+        network = options.hosts_allowed
         if ip == "127.0.0.1":
             self.from_localhost = True
+            allow = True
             logger.info("WebSocket: %s: origin is localhost" % self.name)
-            return True
+        elif network.upper() == "ANY":
+            # Any hosts are allowed
+            allow = True
+            logger.error("WebSocket: %s: any hosts are allowed" % self.name)
+        elif ip_in_network(ip, network):
+            allow = True
+            logger.error("WebSocket: %s: " % self.name +
+                         "client is in the allowed network: %s" % network)
         else:
-            self.from_localhost = False
-            # FIXME/TODO: check whether from local LAN (or in same subnet)??
+            allow = False
             logger.error("WebSocket: %s: " % self.name +
-                         "ONLY allow access from localhost at the moment :(")
-            return False
+                         "client is NOT in the allowed network: %s" % network)
+        return allow
 
     def open(self):
         """Invoked when a new WebSocket is opened by the client."""
-        logger.info("WebSocket: %s: opened" % self.name)
+        logger.info("WebSocket: {0}: opened".format(self.name))
+        logger.info("Allowed hosts: {0}".format(options.hosts_allowed))
 
     def on_close(self):
         """Invoked when a new WebSocket is closed by the client."""
-- 
cgit v1.2.2