From 52beba86952c37a07804fbf7d7468b23650cfc4a Mon Sep 17 00:00:00 2001 From: Joel Martin Date: Thu, 12 Jul 2012 19:10:12 -0500 Subject: [PATCH] Move target-list processing to websockify. websocket.py has no concept of target/proxy so any target processing should happen in websockify itself. Also: - remove URL parsing imports from websocket.py since they are not needed with SimpleHTTPRequestHandler doing the parsing. - read the absolute path of the target_list file on startup since the --web option will change directories if set. --- websocket.py | 57 ----------------------------------------------- websockify | 63 +++++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 58 deletions(-) diff --git a/websocket.py b/websocket.py index 937b861..328782a 100644 --- a/websocket.py +++ b/websocket.py @@ -18,7 +18,6 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates import os, sys, time, errno, signal, socket, traceback, select import array, struct -from cgi import parse_qsl from base64 import b64encode, b64decode # Imports that vary by python version @@ -36,8 +35,6 @@ 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, parse_qs, urlparse -except: from urlparse import urlsplit, parse_qs, urlparse # python 2.6 differences try: from hashlib import md5, sha1 @@ -654,10 +651,6 @@ 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(',') @@ -808,56 +801,6 @@ 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") diff --git a/websockify b/websockify index fc06ec0..f19b61e 100755 --- a/websockify +++ b/websockify @@ -14,6 +14,8 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates import socket, optparse, time, os, sys, subprocess from select import select import websocket +try: from urllib.parse import parse_qs, urlparse +except: from urlparse import parse_qs, urlparse class WebSocketProxy(websocket.WebSocketServer): """ @@ -75,6 +77,9 @@ Traffic Legend: "REBIND_OLD_PORT": str(kwargs['listen_port']), "REBIND_NEW_PORT": str(self.target_port)}) + if self.target_list: + self.target_list = os.path.abspath(self.target_list) + websocket.WebSocketServer.__init__(self, *args, **kwargs) def run_wrap_cmd(self): @@ -153,6 +158,11 @@ Traffic Legend: """ Called after a new WebSocket connection has been established. """ + # Checks if we receive a token, and look + # for a valid target for it then + if self.target_list: + (self.target_host, self.target_port) = self.get_target(self.target_list, self.path) + # Connect to the target if self.wrap_cmd: msg = "connecting to command: %s" % (" ".join(self.wrap_cmd), self.target_port) @@ -183,6 +193,55 @@ Traffic Legend: self.target_host, self.target_port)) raise + def get_target(self, target_list, path): + """ + Parses the path, extracts a token, and looks for a valid + target for that token in the configuration file(s). Sets + target_host and target_port if successful + """ + # 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(target_list): + folder = target_list + targets = os.listdir(target_list) + # If its a file, just add it to the list + else: + folder = os.path.dirname(target_list) + targets = [os.path.basename(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 + found = True + break + + if not found: + raise self.EClose("Token check failed") + + return host, port + def do_proxy(self, target): """ Proxy client WebSocket to normal target socket. @@ -275,7 +334,9 @@ def websockify_init(): 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") + 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