From ed889b238e67d947a203c908de2d94e4d72f0098 Mon Sep 17 00:00:00 2001 From: james-portman Date: Fri, 13 Feb 2015 22:47:17 +0000 Subject: [PATCH] adding get target via api allows to get targets from local rest/json API rather than hardcoded files --- websockify/websocketproxy.py | 49 ++++++++++++++++++++++++++++++------ 1 file changed, 42 insertions(+), 7 deletions(-) diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py index ac9ed633..a6b0499b 100755 --- a/websockify/websocketproxy.py +++ b/websockify/websocketproxy.py @@ -11,7 +11,7 @@ ''' -import signal, socket, optparse, time, os, sys, subprocess, logging +import signal, socket, optparse, time, os, sys, subprocess, logging, urllib2, json try: from socketserver import ForkingMixIn except: from SocketServer import ForkingMixIn try: from http.server import HTTPServer @@ -46,6 +46,8 @@ def new_websocket_client(self): # for a valid target for it then if self.server.target_cfg: (self.server.target_host, self.server.target_port) = self.get_target(self.server.target_cfg, self.path) + elif self.server.target_api: + (self.server.target_host, self.server.target_port) = self.get_target_via_api(self.server.target_api, self.path) # Connect to the target if self.server.wrap_cmd: @@ -73,7 +75,7 @@ def new_websocket_client(self): if tsock: tsock.shutdown(socket.SHUT_RDWR) tsock.close() - if self.verbose: + if self.verbose: self.log_message("%s:%s: Closed target", self.server.target_host, self.server.target_port) raise @@ -117,6 +119,28 @@ def get_target(self, target_cfg, path): else: raise self.EClose("Token '%s' not found" % token) + def get_target_via_api(self,target_api,path): + """ + Parses the path, extracts a token, and looks for a valid + target for that token in the Installer API. Sets + target_host and target_port if successful + """ + + # Extract the token parameter from url + args = parse_qs(urlparse(path)[4]) # 4 is the query from url + + if not 'token' in args or not len(args['token']): + raise self.EClose("Token not present") + + # installId to be passed in in request/query string + token = args['token'][0].rstrip('\n') + try: + output = urllib2.urlopen(target_api+token).read() + decoded = json.loads(output) + return decoded['ip'],decoded['port'] + except: + raise self.EClose("Token '%s' not found" % token) + def do_proxy(self, target): """ Proxy client WebSocket to normal target socket. @@ -147,7 +171,7 @@ def do_proxy(self, target): if closed: # TODO: What about blocking on client socket? - if self.verbose: + if self.verbose: self.log_message("%s:%s: Client closed connection", self.server.target_host, self.server.target_port) raise self.CClose(closed['code'], closed['reason']) @@ -196,6 +220,7 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, *args, **kwargs): self.unix_target = kwargs.pop('unix_target', None) self.ssl_target = kwargs.pop('ssl_target', None) self.target_cfg = kwargs.pop('target_cfg', None) + self.target_api = kwargs.pop('target_api', None) # Last 3 timestamps command was run self.wrap_times = [0, 0, 0] @@ -254,6 +279,9 @@ def started(self): if self.target_cfg: msg = " - proxying from %s:%s to targets in %s" % ( self.listen_host, self.listen_port, self.target_cfg) + elif self.target_api: + msg = " - proxying from %s:%s to targets (if found) in API" % ( + self.listen_host, self.listen_port) else: msg = " - proxying from %s:%s to %s" % ( self.listen_host, self.listen_port, dst_string) @@ -352,6 +380,11 @@ 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-api", metavar="FILE", + dest="target_api", + help="JSON Rest API address to check token, token is appended" + "e.g. http://localhost/api/[token]" + "API should return ip and port in JSON response") parser.add_option("--target-config", metavar="FILE", dest="target_cfg", help="Configuration file containing valid targets " @@ -365,7 +398,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_cfg or opts.unix_target or opts.target_api): parser.error("Too few arguments") if sys.argv.count('--'): opts.wrap_cmd = args[1:] @@ -390,7 +423,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_cfg or opts.target_api: opts.target_host = None opts.target_port = None else: @@ -434,7 +467,9 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, **kwargs): self.unix_target = kwargs.pop('unix_target', None) self.ssl_target = kwargs.pop('ssl_target', None) self.target_cfg = kwargs.pop('target_cfg', None) + self.target_api = kwargs.pop('target_api', None) self.daemon = False + self.target_api = None self.target_cfg = None # Server configuration @@ -456,8 +491,8 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, **kwargs): if web: os.chdir(web) - - HTTPServer.__init__(self, (listen_host, listen_port), + + HTTPServer.__init__(self, (listen_host, listen_port), RequestHandlerClass)