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.
This commit is contained in:
Joel Martin
parent
e17e1158d8
commit
52beba8695
57
websocket.py
57
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")
|
||||
|
|
|
|||
63
websockify
63
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
|
||||
|
|
|
|||
Loading…
Reference in New Issue