From 87bdd0004717b2bd22ddb9f0a9624a7d13fb4344 Mon Sep 17 00:00:00 2001
From: Radek Podgorny <radek@podgorny.cz>
Date: Tue, 30 Sep 2014 13:41:46 +0200
Subject: [PATCH] allow for specification of host and port directly from path

---
 utils/websockify | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/utils/websockify b/utils/websockify
index 7b3ec111..8ea97052 100755
--- a/utils/websockify
+++ b/utils/websockify
@@ -42,9 +42,12 @@ Traffic Legend:
         """
         Called after a new WebSocket connection has been established.
         """
+        # Gets the target host and port directly from path
+        if self.server.target_direct:
+            (self.server.target_host, self.server.target_port) = self.get_target_direct(self.path)
         # Checks if we receive a token, and look
         # for a valid target for it then
-        if self.server.target_cfg:
+        elif self.server.target_cfg:
             (self.server.target_host, self.server.target_port) = self.get_target(self.server.target_cfg, self.path)
 
         # Connect to the target
@@ -78,6 +81,19 @@ Traffic Legend:
                             self.server.target_host, self.server.target_port)
             raise
 
+    def get_target_direct(self, path):
+        args = parse_qs(urlparse(path)[4])  # 4 is the query from url
+
+        if not 'host' in args or not args['host']:
+            raise self.EClose("Host not present")
+        if not 'port' in args or not args['port']:
+            raise self.EClose("Port not present")
+
+        host = args['host'][0].rstrip('\n')
+        port = args['port'][0].rstrip('\n')
+
+        return host, port
+
     def get_target(self, target_cfg, path):
         """
         Parses the path, extracts a token, and looks for a valid
@@ -195,6 +211,7 @@ class WebSocketProxy(websocket.WebSocketServer):
         self.wrap_mode      = kwargs.pop('wrap_mode', None)
         self.unix_target    = kwargs.pop('unix_target', None)
         self.ssl_target     = kwargs.pop('ssl_target', None)
+        self.target_direct  = kwargs.pop('target_direct', None)
         self.target_cfg     = kwargs.pop('target_cfg', None)
         # Last 3 timestamps command was run
         self.wrap_times    = [0, 0, 0]
@@ -352,6 +369,9 @@ def websockify_init():
     parser.add_option("--prefer-ipv6", "-6",
             action="store_true", dest="source_is_ipv6",
             help="prefer IPv6 when resolving source_addr")
+    parser.add_option("--target-direct", action="store_true",
+            dest="target_direct",
+            help="Accept target host and port directly from path")
     parser.add_option("--target-config", metavar="FILE",
             dest="target_cfg",
             help="Configuration file containing valid targets "
@@ -365,7 +385,7 @@ def websockify_init():
         logging.getLogger(WebSocketProxy.log_prefix).setLevel(logging.DEBUG)
 
     # Sanity checks
-    if len(args) < 2 and not (opts.target_cfg or opts.unix_target):
+    if len(args) < 2 and not (opts.target_direct or opts.target_cfg or opts.unix_target):
         parser.error("Too few arguments")
     if sys.argv.count('--'):
         opts.wrap_cmd = args[1:]
@@ -390,7 +410,7 @@ def websockify_init():
     try:    opts.listen_port = int(opts.listen_port)
     except: parser.error("Error parsing listen port")
 
-    if opts.wrap_cmd or opts.unix_target or opts.target_cfg:
+    if opts.wrap_cmd or opts.unix_target or opts.target_direct or opts.target_cfg:
         opts.target_host = None
         opts.target_port = None
     else:
@@ -433,6 +453,7 @@ class LibProxyServer(ForkingMixIn, HTTPServer):
         self.wrap_mode      = kwargs.pop('wrap_mode', None)
         self.unix_target    = kwargs.pop('unix_target', None)
         self.ssl_target     = kwargs.pop('ssl_target', None)
+        self.target_direct  = kwargs.pop('target_direct', None)
         self.target_cfg     = kwargs.pop('target_cfg', None)
         self.daemon = False
         self.target_cfg = None