From ca549cf081b6b1ea3bcfd933552eba527afc1722 Mon Sep 17 00:00:00 2001
From: Andrew Morgan <1342360+anoadragon453@users.noreply.github.com>
Date: Tue, 18 Dec 2018 14:36:14 +0100
Subject: [PATCH] Allow IP addresses via WebSocket API (#1819)

* Allow IP addresses via WebSocket API

* Switch to socket. Add host
---
 src/Ui/UiRequest.py | 39 +++++++++++++++++++++++++++++++++++++--
 src/Ui/UiServer.py  |  6 ++++--
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/src/Ui/UiRequest.py b/src/Ui/UiRequest.py
index ade3f4f2..42ca72b8 100644
--- a/src/Ui/UiRequest.py
+++ b/src/Ui/UiRequest.py
@@ -4,6 +4,7 @@ import os
 import mimetypes
 import json
 import cgi
+import socket
 
 import gevent
 
@@ -48,10 +49,45 @@ class UiRequest(object):
         self.user = None
         self.script_nonce = None  # Nonce for script tags in wrapper html
 
+    # Test if a string is a valid IP address
+    def isIp(self, host, strip_port=False):
+        if strip_port:
+            # Remove the port from the IP address
+            host = ":".join(host.split(":")[:-1])
+
+        try:
+            # This function will return an exception on a non-valid IP
+            # address
+            socket.inet_aton(host)
+            return True
+
+        except socket.error:
+            # Try for a IPv6 address
+            try:
+                socket.inet_pton(socket.AF_INET6, host)
+                return True
+
+            except socket.error:
+                if not strip_port:
+                    # Try stripping the port and re-checking
+                    return self.isIp(host, strip_port=True)
+
+        return False
+
+    def learnHost(self, host):
+        self.server.allowed_hosts.add(host)
+        self.server.log.info("Added %s as allowed host" % host)
+
     def isHostAllowed(self, host):
         if host in self.server.allowed_hosts:
             return True
 
+        # Allow any IP address as they are not affected by DNS rebinding
+        # attacks
+        if self.isIp(host):
+            self.learnHost(host)
+            return True
+
         if self.isProxyRequest():  # Support for chrome extension proxy
             if self.server.site_manager.isDomain(host):
                 return True
@@ -61,8 +97,7 @@ class UiRequest(object):
         if self.server.learn_allowed_host:
             # Learn the first request's host as allowed one
             self.server.learn_allowed_host = False
-            self.server.allowed_hosts.add(host)
-            self.server.log.info("Added %s as allowed host" % host)
+            self.learnHost(host)
             return True
 
         return False
diff --git a/src/Ui/UiServer.py b/src/Ui/UiServer.py
index 9026c4e8..c82f0280 100644
--- a/src/Ui/UiServer.py
+++ b/src/Ui/UiServer.py
@@ -62,7 +62,9 @@ class UiServer:
             self.allowed_hosts = set(config.ui_host)
             self.learn_allowed_host = False
         elif config.ui_ip == "127.0.0.1":
-            self.allowed_hosts = set(["zero", "localhost:%s" % config.ui_port, "127.0.0.1:%s" % config.ui_port])
+            # IP Addresses are inherently allowed as they are immune to DNS
+            # rebinding attacks.
+            self.allowed_hosts = set(["zero", "localhost:%s" % config.ui_port])
             # "URI producers and normalizers should omit the port component and
             # its ':' delimiter if port is empty or if its value would be the
             # same as that of the scheme's default."
@@ -70,7 +72,7 @@ class UiServer:
             # As a result, we need to support portless hosts if port 80 is in
             # use.
             if config.ui_port == 80:
-                self.allowed_hosts.update(["localhost", "127.0.0.1"])
+                self.allowed_hosts.update(["localhost"])
             self.learn_allowed_host = False
         else:
             self.allowed_hosts = set([])