webui: Add "utils.py" to get the IP address

FIXME/TODO:
How to determine the WebSocket origin is in the same subnet as the
server?  An additional network mask required to determine this.
How does this additional mask passed?

2 files changed, 85 insertions, 11 deletions

diff --git a/fg21sim/webui/utils.py b/fg21sim/webui/utils.py
new file mode 100644
index 0000000..1cea69f
--- /dev/null
+++ b/fg21sim/webui/utils.py
@@ -0,0 +1,72 @@
+# Copyright (c) 2016 Weitian LI <liweitianux@live.com>
+# MIT license
+
+"""
+Utilities for the Web UI
+"""
+
+
+from urllib.parse import urlparse
+import socket
+
+
+def get_host_ip(url):
+    """
+    This function parses the input URL to get the hostname (or an IP),
+    then the hostname is further resolved to its IP address.
+
+    Parameters
+    ----------
+    url : str
+        An URL string, which generally has the following format:
+        ``scheme://netloc/path;parameters?query#fragment``
+        while the ``netloc`` may look like ``user:pass@example.com:port``
+
+    Returns
+    -------
+    ip : str
+        An IPv4 address string.
+        If something wrong happens (e.g., ``gaierror``), then ``None``
+        is returned.
+    """
+    netloc = urlparse(url).netloc
+    hostname = netloc.split("@")[-1].split(":")[0]
+    try:
+        ip = socket.gethostbyname(hostname)
+    except socket.gaierror:
+        ip = None
+    return ip
+
+
+def get_local_ip(host="localhost", timeout=3.0):
+    """
+    Get the local IP address of this machine where this script runs.
+
+    A dummy socket will be created and connects to the given host, then
+    the valid local IP address used in this connection can be obtained.
+
+    Parameters
+    ----------
+    host : str
+        The host to which will be connected by a dummy socket, in order
+        to determine the valid IP address.
+    timeout : float
+        Timeout (in seconds) on the blocking socket operations (e.g.,
+        ``connect()``)
+
+    Returns
+    -------
+    ip : str
+        The local IPv4 address of this machine as a string.
+        If something wrong happens (e.g., ``gaierror``), then ``None``
+        is returned.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+        try:
+            s.settimeout(timeout)
+            # Use 0 as the port will let OS determine the free port
+            s.connect((host, 0))
+            ip = s.getsockname()[0]
+        except (socket.gaierror, socket.timeout):
+            ip = None
+    return ip
diff --git a/fg21sim/webui/websocket.py b/fg21sim/webui/websocket.py
index 5b9d5ea..0677f83 100644
--- a/fg21sim/webui/websocket.py
+++ b/fg21sim/webui/websocket.py
@@ -17,11 +17,12 @@ References
   http://caniuse.com/#feat=websockets
 """
 
-from urllib.parse import urlparse
 import logging
 
 import tornado.websocket
 
+from .utils import get_host_ip
+
 
 logger = logging.getLogger(__name__)
 
@@ -63,6 +64,7 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
 
     def on_close(self):
         """Invoked when a new WebSocket is closed by the client."""
+        code, reason = None, None
         if hasattr(self, "close_code"):
             code = self.close_code
         if hasattr(self, "close_reason"):
@@ -84,15 +86,15 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
         Currently, only allow access from the ``localhost``
         (i.e., 127.0.0.1) and local LAN.
         """
-        logger.info("WebSocket: {0}: access origin: {1}".format(
-            self.name, origin))
-        # NOTE: `urlparse`: `scheme://netloc/path;parameters?query#fragment`
-        netloc = urlparse(origin).netloc
-        # NOTE: `netloc` example: `user:pass@example.com:8080`
-        host = netloc.split("@")[-1].split(":")[0]
-        if host in ["localhost", "127.0.0.1"]:
+        logger.info("WebSocket: {0}: origin: {1}".format(self.name, origin))
+        ip = get_host_ip(url=origin)
+        if ip == "127.0.0.1":
             self.from_localhost = True
-            logger.info("WebSocket: %s: access from localhost" % self.name)
+            logger.info("WebSocket: %s: origin is localhost" % self.name)
             return True
-        # XXX/TODO: check whether from the local LAN ??
-        return False
+        else:
+            self.from_localhost = False
+            # FIXME/TODO: check whether from local LAN (or in same subnet)??
+            logger.error("WebSocket: %s: " % self.name +
+                         "ONLY allow access from localhost at the moment :(")
+            return False