diff --git a/utils/websocket.py b/utils/websocket.py index 8194914a..37b69c86 100644 --- a/utils/websocket.py +++ b/utils/websocket.py @@ -16,7 +16,8 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates ''' -import os, sys, time, errno, signal, socket, struct, traceback, select +import os, sys, time, errno, signal, socket, traceback, select +import struct, array from cgi import parse_qsl from base64 import b64encode, b64decode @@ -28,6 +29,7 @@ if sys.hexversion > 0x3000000: from urllib.parse import urlsplit b2s = lambda buf: buf.decode('latin_1') s2b = lambda s: s.encode('latin_1') + s2a = lambda s: s else: # python 2.X from cStringIO import StringIO @@ -36,6 +38,7 @@ else: # No-ops b2s = lambda buf: buf s2b = lambda s: s + s2a = lambda s: [ord(c) for c in s] if sys.hexversion >= 0x2060000: # python >= 2.6 @@ -54,7 +57,7 @@ for mod, sup in [('numpy', 'HyBi protocol'), globals()[mod] = __import__(mod) except ImportError: globals()[mod] = None - print("WARNING: no '%s' module, %s support disabled" % ( + print("WARNING: no '%s' module, %s decode may be slower" % ( mod, sup)) @@ -88,7 +91,8 @@ Sec-WebSocket-Accept: %s\r def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False, verbose=False, cert='', key='', ssl_only=None, - daemon=False, record='', web=''): + daemon=False, record='', web='', + run_once=False, timeout=0): # settings self.verbose = verbose @@ -96,6 +100,11 @@ Sec-WebSocket-Accept: %s\r self.listen_port = listen_port self.ssl_only = ssl_only self.daemon = daemon + self.run_once = run_once + self.timeout = timeout + + self.launch_time = time.time() + self.ws_connection = False self.handler_id = 1 # Make paths settings absolute @@ -207,6 +216,38 @@ Sec-WebSocket-Accept: %s\r os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdout.fileno()) os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno()) + @staticmethod + def unmask(buf, f): + pstart = f['hlen'] + 4 + pend = pstart + f['length'] + if numpy: + b = c = s2b('') + if f['length'] >= 4: + mask = numpy.frombuffer(buf, dtype=numpy.dtype('= 4: - mask = numpy.frombuffer(buf, dtype=numpy.dtype('= 2.6 and numpy module is required for HyBi-07 or greater") - # HyBi-07 report version 7 # HyBi-08 - HyBi-12 report version 8 # HyBi-13 reports version 13 @@ -669,6 +690,9 @@ Sec-WebSocket-Accept: %s\r self.msg("%s: %s WebSocket connection" % (address[0], stype)) self.msg("%s: Version %s, base64: '%s'" % (address[0], self.version, self.base64)) + if self.path != '/': + self.msg("%s: Path: '%s'" % (address[0], self.path)) + # Send server WebSockets handshake response #self.msg("sending response [%s]" % response) @@ -727,6 +751,7 @@ Sec-WebSocket-Accept: %s\r self.rec = open(fname, 'w+') self.rec.write("var VNC_frame_data = [\n") + self.ws_connection = True self.new_client() except self.EClose: _, exc, _ = sys.exc_info() @@ -777,6 +802,12 @@ Sec-WebSocket-Accept: %s\r startsock = None pid = err = 0 + time_elapsed = time.time() - self.launch_time + if self.timeout and time_elapsed > self.timeout: + self.msg('listener exit due to --timeout %s' + % self.timeout) + break + try: self.poll() @@ -799,7 +830,14 @@ Sec-WebSocket-Accept: %s\r else: raise - if Process: + if self.run_once: + # Run in same process if run_once + self.top_new_client(startsock, address) + if self.ws_connection : + self.msg('%s: exiting due to --run-once' + % address[0]) + break + elif Process: self.vmsg('%s: new handler Process' % address[0]) p = Process(target=self.top_new_client, args=(startsock, address)) diff --git a/utils/websockify b/utils/websockify index 59b6e0de..4521a963 100755 --- a/utils/websockify +++ b/utils/websockify @@ -226,6 +226,10 @@ if __name__ == '__main__': 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,