#!/usr/bin/env python
'''
A WebSocket to TCP socket proxy with support for "wss://" encryption.
Copyright 2011 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.
"""
# Connect to the target
self.msg("connecting to: %s:%s" % (
self.target_host, self.target_port))
tsock = self.socket(self.target_host, self.target_port,
connect=True)
if self.verbose and not self.daemon:
print(self.traffic_legend)
# Start proxying
try:
self.do_proxy(tsock)
except:
if tsock:
tsock.shutdown(socket.SHUT_RDWR)
tsock.close()
self.vmsg("%s:%s: Target closed" %(
self.target_host, self.target_port))
raise
def do_proxy(self, target):
"""
Proxy client WebSocket to normal target socket.
"""
cqueue = []
c_pend = 0
tqueue = []
rlist = [self.client, target]
while True:
wlist = []
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.client in ins:
# Receive client data, decode it, and queue for target
bufs, closed = self.recv_frames()
tqueue.extend(bufs)
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("--run-once", action="store_true",
help="handle a single WebSocket connection and exit")
parser.add_option("--timeout", type=int, default=0,
help="after TIMEOUT seconds exit when not connected")
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].rsplit(':', 1)
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].rsplit(':', 1)
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()