diff --git a/websocket.py b/websocket.py index 03bedff..937b861 100644 --- a/websocket.py +++ b/websocket.py @@ -36,8 +36,8 @@ try: from io import StringIO except: from cStringIO import StringIO try: from http.server import SimpleHTTPRequestHandler except: from SimpleHTTPServer import SimpleHTTPRequestHandler -try: from urllib.parse import urlsplit -except: from urlparse import urlsplit +try: from urllib.parse import urlsplit, parse_qs, urlparse +except: from urlparse import urlsplit, parse_qs, urlparse # python 2.6 differences try: from hashlib import md5, sha1 @@ -654,6 +654,10 @@ Sec-WebSocket-Accept: %s\r h = self.headers = wsh.headers path = self.path = wsh.path + # Checks if we receive a token, and look + # for a valid target for it then + self.set_target(path) + prot = 'WebSocket-Protocol' protocols = h.get('Sec-'+prot, h.get(prot, '')).split(',') @@ -804,6 +808,56 @@ Sec-WebSocket-Accept: %s\r # Original socket closed by caller self.client.close() + def set_target(self, path): + """ + Parses the path, extracts a token, and looks for a valid + target for that token in the configuration files. Sets + target_host and port if successful + """ + if self.target_list: + # The files in targets contain the lines + # in the form of host:port:token + + # Extract the token parameter from url + args = parse_qs(urlparse(path)[4]) # 4 is the query from url + + if not len(args['token']): + raise self.EClose("Token not present") + + token = args['token'][0].rstrip('\n') + + # If target list is a directory, then list all files in it + if os.path.isdir(self.target_list): + folder = self.target_list + targets = os.listdir(self.target_list) + # If its a file, just add it to the list + else: + folder = os.path.dirname(self.target_list) + targets = [os.path.basename(self.target_list)] + + target_lines = [] + + try: + # extract lines from every config file + for filename in targets: + f = open(folder + '/' + filename, 'r') + target_lines.extend(f.readlines()) + except: + raise self.EClose("Could not read token file(s)") + + # search for the line matching the provided token + found = False + for target in target_lines: + host, port, file_token = target.rstrip('\n').split(':') + if file_token == token: # found, set target + self.target_host = host + self.target_port = port + found = True + break + + if not found: + raise self.EClose("Token check failed") + def new_client(self): """ Do something with a WebSockets client connection. """ raise("WebSocketServer.new_client() must be overloaded") @@ -941,4 +995,3 @@ class WSRequestHandler(SimpleHTTPRequestHandler): def log_message(self, f, *args): # Save instead of printing self.last_message = f % args - diff --git a/websockify b/websockify index 761fe90..fc06ec0 100755 --- a/websockify +++ b/websockify @@ -45,6 +45,7 @@ Traffic Legend: self.wrap_mode = kwargs.pop('wrap_mode') self.unix_target = kwargs.pop('unix_target') self.ssl_target = kwargs.pop('ssl_target') + self.target_list = kwargs.pop('target_list') # Last 3 timestamps command was run self.wrap_times = [0, 0, 0] @@ -97,7 +98,11 @@ Traffic Legend: else: dst_string = "%s:%s" % (self.target_host, self.target_port) - msg = " - proxying from %s:%s to %s" % ( + if self.target_list: + msg = " - proxying from %s:%s to targets in %s" % ( + self.listen_host, self.listen_port, self.target_list) + else: + msg = " - proxying from %s:%s to %s" % ( self.listen_host, self.listen_port, dst_string) if self.ssl_target: @@ -235,7 +240,7 @@ Traffic Legend: def websockify_init(): usage = "\n %prog [options]" - usage += " [source_addr:]source_port target_addr:target_port" + usage += " [source_addr:]source_port [target_addr:target_port]" usage += "\n %prog [options]" usage += " [source_addr:]source_port -- WRAP_COMMAND_LINE" parser = optparse.OptionParser(usage=usage) @@ -269,10 +274,12 @@ 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-list", metavar="FILE", + help="Configuration file containing valid targets in the form host:port:token or, alternatively, a directory containing configuration files of this form") (opts, args) = parser.parse_args() # Sanity checks - if len(args) < 1: + if len(args) < 2 and not opts.target_list: parser.error("Too few arguments") if sys.argv.count('--'): opts.wrap_cmd = args[1:] @@ -297,7 +304,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: + if opts.wrap_cmd or opts.unix_target or opts.target_list: opts.target_host = None opts.target_port = None else: