From df1aa028da122acd5ea4474182dfda567d1c839a Mon Sep 17 00:00:00 2001
From: Luca Capacci <luca.capacci@cryptonet.it>
Date: Mon, 28 Sep 2015 13:00:55 +0200
Subject: [PATCH] Added --auth-host-port option

---
 websockify/websocketproxy.py | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py
index 46ab545..3651556 100755
--- a/websockify/websocketproxy.py
+++ b/websockify/websocketproxy.py
@@ -53,9 +53,16 @@ Traffic Legend:
 
         if self.server.auth_plugin:
             try:
-                self.server.auth_plugin.authenticate(
-                    headers=self.headers, target_host=self.server.target_host,
-                    target_port=self.server.target_port)
+                if self.server.auth_host_port:
+                    server_target_host, server_target_port = self.server.auth_plugin.authenticate(headers=self.headers,
+                                                                                                  target_host=self.server.target_host,
+                                                                                                  target_port=self.server.target_port)
+                    self.server.target_host = server_target_host
+                    self.server.target_port = server_target_port
+                else:
+                    self.server.auth_plugin.authenticate(
+                        headers=self.headers, target_host=self.server.target_host,
+                        target_port=self.server.target_port)
             except auth.AuthenticationError:
                 ex = sys.exc_info()[1]
                 self.send_auth_error(ex)
@@ -229,6 +236,7 @@ class WebSocketProxy(websocket.WebSocketServer):
 
         self.token_plugin = kwargs.pop('token_plugin', None)
         self.auth_plugin = kwargs.pop('auth_plugin', None)
+        self.auth_host_port = kwargs.pop('auth_host_port', False)
 
         # Last 3 timestamps command was run
         self.wrap_times    = [0, 0, 0]
@@ -288,6 +296,9 @@ class WebSocketProxy(websocket.WebSocketServer):
         if self.token_plugin:
             msg = "  - proxying from %s:%s to targets generated by %s" % (
                 self.listen_host, self.listen_port, type(self.token_plugin).__name__)
+        elif self.auth_host_port:
+            msg = "  - proxying from %s:%s to targets generated by %s" % (
+                self.listen_host, self.listen_port, type(self.auth_plugin).__name__)
         else:
             msg = "  - proxying from %s:%s to %s" % (
                 self.listen_host, self.listen_port, dst_string)
@@ -407,6 +418,8 @@ def websockify_init():
     parser.add_option("--auth-source", default=None, metavar="ARG",
                       help="an argument to be passed to the auth plugin"
                            "on instantiation")
+    parser.add_option("--auth-host-port", action="store_true",
+                      help="let the auth plugin set host and port")
     parser.add_option("--auto-pong", action="store_true",
             help="Automatically respond to ping frames with a pong")
     parser.add_option("--heartbeat", type=int, default=0,
@@ -423,6 +436,8 @@ def websockify_init():
     if opts.auth_source and not opts.auth_plugin:
         parser.error("You must use --auth-plugin to use --auth-source")
 
+    if opts.auth_host_port and not opts.auth_plugin:
+        parser.error("You must use --auth-plugin to use --auth-host-port")
 
     # Transform to absolute path as daemon may chdir
     if opts.target_cfg:
@@ -435,7 +450,7 @@ def websockify_init():
     del opts.target_cfg
 
     # Sanity checks
-    if len(args) < 2 and not (opts.token_plugin or opts.unix_target):
+    if len(args) < 2 and not (opts.token_plugin or opts.unix_target or opts.auth_host_port):
         parser.error("Too few arguments")
     if sys.argv.count('--'):
         opts.wrap_cmd = args[1:]
@@ -460,7 +475,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.token_plugin:
+    if opts.wrap_cmd or opts.unix_target or opts.token_plugin or opts.auth_host_port:
         opts.target_host = None
         opts.target_port = None
     else: