Joel Martin
88b71ce171
- Add initial IETF-07 (HyBi-07) protocol version support. This version still uses base64 encoding since the API for binary support is not yet finalized. - Move socket send and recieve functions into the WebSocketServer class instead of having the sub-class do this. This simplifies sub-classes somewhat. The send_frame routine now returns the number of frames that were unable to be sent. If this value is non-zero then the sub-class should call again when the socket is ready until the pending frames count is 0. - Do traffic reporting in the main class instead. - When the client is HyBi style (i.e. IETF-07) then use the sub-protocol header to select whether to do base64 encoding or simply send the frame data raw (binary). Update include/websock.js to send a 'base64' protocol selector. Once the API support binary, then the client will need to detect this and set the protocol to 'binary'.
299 lines
10 KiB
Python
Executable File
299 lines
10 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
'''
|
|
A WebSocket to TCP socket proxy with support for "wss://" encryption.
|
|
Copyright 2010 Joel Martin
|
|
Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
|
|
|
|
You can make a cert/key with openssl using:
|
|
openssl req -new -x509 -days 365 -nodes -out self.pem -keyout self.pem
|
|
as taken from http://docs.python.org/dev/library/ssl.html#certificates
|
|
|
|
'''
|
|
|
|
import socket, optparse, time, os, sys, subprocess
|
|
from select import select
|
|
from websocket import WebSocketServer
|
|
|
|
class WebSocketProxy(WebSocketServer):
|
|
"""
|
|
Proxy traffic to and from a WebSockets client to a normal TCP
|
|
socket server target. All traffic to/from the client is base64
|
|
encoded/decoded to allow binary data to be sent/received to/from
|
|
the target.
|
|
"""
|
|
|
|
buffer_size = 65536
|
|
|
|
traffic_legend = """
|
|
Traffic Legend:
|
|
} - Client receive
|
|
}. - Client receive partial
|
|
{ - Target receive
|
|
|
|
> - Target send
|
|
>. - Target send partial
|
|
< - Client send
|
|
<. - Client send partial
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
# Save off proxy specific options
|
|
self.target_host = kwargs.pop('target_host')
|
|
self.target_port = kwargs.pop('target_port')
|
|
self.wrap_cmd = kwargs.pop('wrap_cmd')
|
|
self.wrap_mode = kwargs.pop('wrap_mode')
|
|
# Last 3 timestamps command was run
|
|
self.wrap_times = [0, 0, 0]
|
|
|
|
if self.wrap_cmd:
|
|
rebinder_path = ['./', os.path.dirname(sys.argv[0])]
|
|
self.rebinder = None
|
|
|
|
for rdir in rebinder_path:
|
|
rpath = os.path.join(rdir, "rebind.so")
|
|
if os.path.exists(rpath):
|
|
self.rebinder = rpath
|
|
break
|
|
|
|
if not self.rebinder:
|
|
raise Exception("rebind.so not found, perhaps you need to run make")
|
|
|
|
self.target_host = "127.0.0.1" # Loopback
|
|
# Find a free high port
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.bind(('', 0))
|
|
self.target_port = sock.getsockname()[1]
|
|
sock.close()
|
|
|
|
os.environ.update({
|
|
"LD_PRELOAD": self.rebinder,
|
|
"REBIND_OLD_PORT": str(kwargs['listen_port']),
|
|
"REBIND_NEW_PORT": str(self.target_port)})
|
|
|
|
WebSocketServer.__init__(self, *args, **kwargs)
|
|
|
|
def run_wrap_cmd(self):
|
|
print "Starting '%s'" % " ".join(self.wrap_cmd)
|
|
self.wrap_times.append(time.time())
|
|
self.wrap_times.pop(0)
|
|
self.cmd = subprocess.Popen(
|
|
self.wrap_cmd, env=os.environ)
|
|
self.spawn_message = True
|
|
|
|
def started(self):
|
|
"""
|
|
Called after Websockets server startup (i.e. after daemonize)
|
|
"""
|
|
# Need to call wrapped command after daemonization so we can
|
|
# know when the wrapped command exits
|
|
if self.wrap_cmd:
|
|
print " - proxying from %s:%s to '%s' (port %s)\n" % (
|
|
self.listen_host, self.listen_port,
|
|
" ".join(self.wrap_cmd), self.target_port)
|
|
self.run_wrap_cmd()
|
|
else:
|
|
print " - proxying from %s:%s to %s:%s\n" % (
|
|
self.listen_host, self.listen_port,
|
|
self.target_host, self.target_port)
|
|
|
|
def poll(self):
|
|
# If we are wrapping a command, check it's status
|
|
|
|
if self.wrap_cmd and self.cmd:
|
|
ret = self.cmd.poll()
|
|
if ret != None:
|
|
self.vmsg("Wrapped command exited (or daemon). Returned %s" % ret)
|
|
self.cmd = None
|
|
|
|
if self.wrap_cmd and self.cmd == None:
|
|
# Response to wrapped command being gone
|
|
if self.wrap_mode == "ignore":
|
|
pass
|
|
elif self.wrap_mode == "exit":
|
|
sys.exit(ret)
|
|
elif self.wrap_mode == "respawn":
|
|
now = time.time()
|
|
avg = sum(self.wrap_times)/len(self.wrap_times)
|
|
if (now - avg) < 10:
|
|
# 3 times in the last 10 seconds
|
|
if self.spawn_message:
|
|
print "Command respawning too fast"
|
|
self.spawn_message = False
|
|
else:
|
|
self.run_wrap_cmd()
|
|
|
|
#
|
|
# Routines above this point are run in the master listener
|
|
# process.
|
|
#
|
|
|
|
#
|
|
# Routines below this point are connection handler routines and
|
|
# will be run in a separate forked process for each connection.
|
|
#
|
|
|
|
def new_client(self):
|
|
"""
|
|
Called after a new WebSocket connection has been established.
|
|
"""
|
|
|
|
self.rec = None
|
|
if self.record:
|
|
# Record raw frame data as a JavaScript compatible file
|
|
fname = "%s.%s" % (self.record,
|
|
self.handler_id)
|
|
self.msg("opening record file: %s" % fname)
|
|
self.rec = open(fname, 'w+')
|
|
self.rec.write("var VNC_frame_data = [\n")
|
|
|
|
# Connect to the target
|
|
self.msg("connecting to: %s:%s" % (
|
|
self.target_host, self.target_port))
|
|
tsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
tsock.connect((self.target_host, self.target_port))
|
|
|
|
if self.verbose and not self.daemon:
|
|
print self.traffic_legend
|
|
|
|
# Start proxying
|
|
try:
|
|
self.do_proxy(tsock)
|
|
except:
|
|
if tsock:
|
|
tsock.close()
|
|
self.vmsg("%s:%s: Target closed" %(
|
|
self.target_host, self.target_port))
|
|
if self.rec:
|
|
self.rec.write("'EOF']\n")
|
|
self.rec.close()
|
|
raise
|
|
|
|
def do_proxy(self, target):
|
|
"""
|
|
Proxy client WebSocket to normal target socket.
|
|
"""
|
|
cqueue = []
|
|
c_pend = 0
|
|
tqueue = []
|
|
rlist = [self.client, target]
|
|
tstart = int(time.time()*1000)
|
|
|
|
while True:
|
|
wlist = []
|
|
tdelta = int(time.time()*1000) - tstart
|
|
|
|
if tqueue: wlist.append(target)
|
|
if cqueue or c_pend: wlist.append(self.client)
|
|
ins, outs, excepts = select(rlist, wlist, [], 1)
|
|
if excepts: raise Exception("Socket exception")
|
|
|
|
if target in outs:
|
|
# Send queued client data to the target
|
|
dat = tqueue.pop(0)
|
|
sent = target.send(dat)
|
|
if sent == len(dat):
|
|
self.traffic(">")
|
|
else:
|
|
# requeue the remaining data
|
|
tqueue.insert(0, dat[sent:])
|
|
self.traffic(".>")
|
|
|
|
|
|
if target in ins:
|
|
# Receive target data, encode it and queue for client
|
|
buf = target.recv(self.buffer_size)
|
|
if len(buf) == 0: raise self.EClose("Target closed")
|
|
|
|
cqueue.append(buf)
|
|
self.traffic("{")
|
|
|
|
|
|
if self.client in outs:
|
|
# Send queued target data to the client
|
|
c_pend = self.send_frames(cqueue)
|
|
cqueue = []
|
|
|
|
#if self.rec:
|
|
# self.rec.write("%s,\n" %
|
|
# repr("{%s{" % tdelta + dat[1:-1]))
|
|
|
|
|
|
if self.client in ins:
|
|
# Receive client data, decode it, and queue for target
|
|
bufs, closed = self.recv_frames()
|
|
tqueue.extend(bufs)
|
|
|
|
#if self.rec:
|
|
# for b in bufs:
|
|
# self.rec.write(
|
|
# repr("}%s}%s" % (tdelta, b)) + ",\n")
|
|
|
|
if closed:
|
|
# TODO: What about blocking on client socket?
|
|
self.send_close()
|
|
raise self.EClose(closed)
|
|
|
|
if __name__ == '__main__':
|
|
usage = "\n %prog [options]"
|
|
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)
|
|
parser.add_option("--verbose", "-v", action="store_true",
|
|
help="verbose messages and per frame traffic")
|
|
parser.add_option("--record",
|
|
help="record sessions to FILE.[session_number]", metavar="FILE")
|
|
parser.add_option("--daemon", "-D",
|
|
dest="daemon", action="store_true",
|
|
help="become a daemon (background process)")
|
|
parser.add_option("--cert", default="self.pem",
|
|
help="SSL certificate file")
|
|
parser.add_option("--key", default=None,
|
|
help="SSL key file (if separate from cert)")
|
|
parser.add_option("--ssl-only", action="store_true",
|
|
help="disallow non-encrypted connections")
|
|
parser.add_option("--web", default=None, metavar="DIR",
|
|
help="run webserver on same port. Serve files from DIR.")
|
|
parser.add_option("--wrap-mode", default="exit", metavar="MODE",
|
|
choices=["exit", "ignore", "respawn"],
|
|
help="action to take when the wrapped program exits "
|
|
"or daemonizes: exit (default), ignore, respawn")
|
|
(opts, args) = parser.parse_args()
|
|
|
|
# Sanity checks
|
|
if len(args) < 2:
|
|
parser.error("Too few arguments")
|
|
if sys.argv.count('--'):
|
|
opts.wrap_cmd = args[1:]
|
|
else:
|
|
opts.wrap_cmd = None
|
|
if len(args) > 2:
|
|
parser.error("Too many arguments")
|
|
|
|
if opts.ssl_only and not os.path.exists(opts.cert):
|
|
parser.error("SSL only and %s not found" % opts.cert)
|
|
|
|
# Parse host:port and convert ports to numbers
|
|
if args[0].count(':') > 0:
|
|
opts.listen_host, opts.listen_port = args[0].split(':')
|
|
else:
|
|
opts.listen_host, opts.listen_port = '', args[0]
|
|
try: opts.listen_port = int(opts.listen_port)
|
|
except: parser.error("Error parsing listen port")
|
|
|
|
if opts.wrap_cmd:
|
|
opts.target_host = None
|
|
opts.target_port = None
|
|
else:
|
|
if args[1].count(':') > 0:
|
|
opts.target_host, opts.target_port = args[1].split(':')
|
|
else:
|
|
parser.error("Error parsing target")
|
|
try: opts.target_port = int(opts.target_port)
|
|
except: parser.error("Error parsing target port")
|
|
|
|
# Create and start the WebSockets proxy
|
|
server = WebSocketProxy(**opts.__dict__)
|
|
server.start_server()
|