Files
noVNC/utils/websockify
Joel Martin 1e50871599 websockify --run-once, --timeout, numpy fallback
Pull websockify 724aa3a.

- Use array module for unmasking HyBi when no numpy module is
    available.

- Detect client close properly when using python 3.

- Print request URL path is specified.

- New option --run-once will exit after handling a single WebSocket
  connection (but not ater flash policy or normal web requests).

- New option --timeout TIME will stop listening for new connections
  after exit after TIME seconds (the master process shuts down).
  Existing WebSocket connections will continue but once all
  connections are closed all processes will terminate.
2011-09-29 16:12:19 -05:00

283 lines
9.6 KiB
Python
Executable File

#!/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()