From af85184e28d8e4333472940bfe1d2eb6436b6733 Mon Sep 17 00:00:00 2001
From: josedpedroso <josedpedroso@gmail.com>
Date: Thu, 5 Jul 2018 00:54:19 +0100
Subject: [PATCH] Added --host-token to allow choosing target by hostname.

---
 tests/test_websocketproxy.py   |  6 +++---
 websockify/websocketproxy.py   | 33 ++++++++++++++++++++++++---------
 websockify/websockifyserver.py |  1 +
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/tests/test_websocketproxy.py b/tests/test_websocketproxy.py
index 81c119e..5bd3a16 100644
--- a/tests/test_websocketproxy.py
+++ b/tests/test_websocketproxy.py
@@ -88,7 +88,7 @@ class ProxyRequestHandlerTestCase(unittest.TestCase):
                 return ("some host", "some port")
 
         host, port = self.handler.get_target(
-            TestPlugin(None), self.handler.path)
+            TestPlugin(None))
 
         self.assertEqual(host, "some host")
         self.assertEqual(port, "some port")
@@ -99,7 +99,7 @@ class ProxyRequestHandlerTestCase(unittest.TestCase):
                 return ("unix_socket", "/tmp/socket")
 
         _, socket = self.handler.get_target(
-            TestPlugin(None), self.handler.path)
+            TestPlugin(None))
 
         self.assertEqual(socket, "/tmp/socket")
 
@@ -109,7 +109,7 @@ class ProxyRequestHandlerTestCase(unittest.TestCase):
                 return None
 
         self.assertRaises(FakeServer.EClose, self.handler.get_target,
-            TestPlugin(None), "https://localhost:6080/websockify?token=blah")
+            TestPlugin(None))
 
     def test_token_plugin(self):
         class TestPlugin(token_plugins.BasePlugin):
diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py
index b95f5a4..b8e3d99 100644
--- a/websockify/websocketproxy.py
+++ b/websockify/websocketproxy.py
@@ -135,23 +135,31 @@ Traffic Legend:
                     self.log_message("%s:%s: Closed target",
                             self.server.target_host, self.server.target_port)
 
-    def get_target(self, target_plugin, path):
+    def get_target(self, target_plugin):
         """
-        Parses the path, extracts a token, and looks up a target
-        for that token using the token plugin. Sets
-        target_host and target_port if successful
+        Gets a token from either the path or the host,
+        depending on --host-token, and looks up a target
+        for that token using the token plugin. Used by
+        validate_connection() to set target_host and target_port.
         """
         # The files in targets contain the lines
         # in the form of token: host:port
 
-        # Extract the token parameter from url
-        args = parse_qs(urlparse(path)[4]) # 4 is the query from url
+        if self.host_token:
+            token = self.headers.get('Host')
 
-        if not 'token' in args or not len(args['token']):
+        else:
+            # Extract the token parameter from url
+            args = parse_qs(urlparse(self.path)[4]) # 4 is the query from url
+
+            if 'token' in args and len(args['token']):
+                token = args['token'][0].rstrip('\n')
+            else:
+                token = None
+
+        if token is None:
             raise self.server.EClose("Token not present")
 
-        token = args['token'][0].rstrip('\n')
-
         result_pair = target_plugin.lookup(token)
 
         if result_pair is not None:
@@ -263,6 +271,7 @@ class WebSocketProxy(websockifyserver.WebSockifyServer):
         self.heartbeat      = kwargs.pop('heartbeat', None)
 
         self.token_plugin = kwargs.pop('token_plugin', None)
+        self.host_token = kwargs.pop('host_token', None)
         self.auth_plugin = kwargs.pop('auth_plugin', None)
 
         # Last 3 timestamps command was run
@@ -451,6 +460,9 @@ def websockify_init():
     parser.add_option("--token-source", default=None, metavar="ARG",
                       help="an argument to be passed to the token plugin "
                            "on instantiation")
+    parser.add_option("--host-token", action="store_true",
+                      help="use the host HTTP header as token instead of the "
+                           "token URL query parameter")
     parser.add_option("--auth-plugin", default=None, metavar="CLASS",
                       help="use a Python class, usually one from websockify.auth_plugins, "
                            "such as BasicHTTPAuth, to determine if a connection is allowed")
@@ -517,6 +529,9 @@ def websockify_init():
     if opts.token_source and not opts.token_plugin:
         parser.error("You must use --token-plugin to use --token-source")
 
+    if opts.host_token and not opts.token_plugin:
+        parser.error("You must use --token-plugin to use --host-token")
+
     if opts.auth_source and not opts.auth_plugin:
         parser.error("You must use --auth-plugin to use --auth-source")
 
diff --git a/websockify/websockifyserver.py b/websockify/websockifyserver.py
index e8c7ae5..b9787a6 100644
--- a/websockify/websockifyserver.py
+++ b/websockify/websockifyserver.py
@@ -93,6 +93,7 @@ class WebSockifyRequestHandler(WebSocketRequestHandler, SimpleHTTPRequestHandler
         self.file_only = getattr(server, "file_only", False)
         self.traffic = getattr(server, "traffic", False)
         self.web_auth = getattr(server, "web_auth", False)
+        self.host_token = getattr(server, "host_token", False)
 
         self.logger = getattr(server, "logger", None)
         if self.logger is None: