joshua
/
websockify
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Python 3.0 support. Use multiprocessing module. 

Multiprocessing:
- Switch to using multiprocessing module for python >= 2.6. For python
  2.4 continue to use the os.fork() method.
- Move the new_client creation into top_new_client method to enable
  multiprocessing refactor.
- Only do SIGCHLD handling for os.fork/python 2.4. When doing our own
  SIGCHLD handling under python 3.0, we can run into a python futex
  hang when reloading a web page rapidly. Multiprocessing does it's
  own child reaping so we only need it with os.fork().

Python 3.0:
- Modify imports to reflect new locations: StringIO from io,
  SimpleHTTPRequestHandler from http.server, urlsplit from
  urllib.parse.
- Convert all print statements to print() calls. This also means no
  comma parameter idiom and only using string formatting.
- Define b2s (bytes-to-string) and s2b (string-to-bytes) which are
  no-ops on python versions prior to python 3. In python 3 these do
  the conversion between string and bytes.
- Use compatible try/except method. Exception variable must be
  extracted using sys.exc_info() rather than as part of the except
  statement.

Python 2.4:
- Now degrades more gracefully if ssl module is not found. It will
  still run, but will refuse SSL connections.
- Doesn't support HyBi-07 version due to numpy and struct.unpack_from
  requirement.
This commit is contained in:
Joel Martin 2011-05-18 11:09:10 -05:00
parent 85f84ef280
commit c8587115bc
2 changed files with 151 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

222
websocket.py
View File

@ -2,7 +2,7 @@
''' '''
Python WebSocket library with support for "wss://" encryption. Python WebSocket library with support for "wss://" encryption.
Copyright 2010 Joel Martin Copyright 2011 Joel Martin
Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3) Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
Supports following protocol versions: Supports following protocol versions:
@ -16,23 +16,49 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
''' '''
import sys, socket, ssl, struct, traceback, select import sys, socket, struct, traceback, select
import os, resource, errno, signal # daemonizing import os, resource, errno, signal # daemonizing
from SimpleHTTPServer import SimpleHTTPRequestHandler from cgi import parse_qsl
from cStringIO import StringIO
from base64 import b64encode, b64decode from base64 import b64encode, b64decode
try:
# Imports that vary by python version
if sys.hexversion > 0x3000000:
# python >= 3.0
from io import StringIO
from http.server import SimpleHTTPRequestHandler
from urllib.parse import urlsplit
b2s = lambda buf: buf.decode('latin_1')
s2b = lambda s: s.encode('latin_1')
else:
# python 2.X
from cStringIO import StringIO
from SimpleHTTPServer import SimpleHTTPRequestHandler
from urlparse import urlsplit
# No-ops
b2s = lambda buf: buf
s2b = lambda s: s
if sys.hexversion >= 0x2060000:
# python >= 2.6
from multiprocessing import Process
from hashlib import md5, sha1 from hashlib import md5, sha1
except: else:
# Support python 2.4 # python < 2.6
Process = None
from md5 import md5 from md5 import md5
from sha import sha as sha1 from sha import sha as sha1
# Degraded functionality if these imports are missing
try: try:
# Required for HyBi encode/decode
import numpy, ctypes import numpy, ctypes
except: except ImportError:
numpy = ctypes = None numpy = ctypes = None
from urlparse import urlsplit
from cgi import parse_qsl try:
import ssl
except:
ssl = None
class WebSocketServer(object): class WebSocketServer(object):
""" """
@ -88,20 +114,25 @@ Sec-WebSocket-Accept: %s\r
self.handler_id = 1 self.handler_id = 1
print "WebSocket server settings:" print("WebSocket server settings:")
print " - Listen on %s:%s" % ( print(" - Listen on %s:%s" % (
self.listen_host, self.listen_port) self.listen_host, self.listen_port))
print " - Flash security policy server" print(" - Flash security policy server")
if self.web: if self.web:
print " - Web server" print(" - Web server")
if os.path.exists(self.cert): if ssl and self.ssl_only:
print " - SSL/TLS support" raise Exception("No 'ssl' module and SSL only specified")
if self.ssl_only: if ssl:
print " - Deny non-SSL/TLS connections" if os.path.exists(self.cert):
print(" - SSL/TLS support")
if self.ssl_only:
print(" - Deny non-SSL/TLS connections")
else:
print(" - No SSL/TLS support (no cert file)")
else: else:
print " - No SSL/TLS support (no cert file)" print(" - No SSL/TLS support (no 'ssl' module)")
if self.daemon: if self.daemon:
print " - Backgrounding (daemon)" print(" - Backgrounding (daemon)")
# #
# WebSocketServer static methods # WebSocketServer static methods
@ -133,7 +164,8 @@ Sec-WebSocket-Accept: %s\r
try: try:
if fd != keepfd: if fd != keepfd:
os.close(fd) os.close(fd)
except OSError, exc: except OSError:
_, exc, _ = sys.exc_info()
if exc.errno != errno.EBADF: raise if exc.errno != errno.EBADF: raise
# Redirect I/O to /dev/null # Redirect I/O to /dev/null
@ -164,7 +196,7 @@ Sec-WebSocket-Accept: %s\r
elif payload_len >= 65536: elif payload_len >= 65536:
header = struct.pack('>BBQ', b1, 127, payload_len) header = struct.pack('>BBQ', b1, 127, payload_len)
#print "Encoded: %s" % repr(header + buf) #print("Encoded: %s" % repr(header + buf))
return header + buf return header + buf
@ -238,7 +270,7 @@ Sec-WebSocket-Accept: %s\r
b = numpy.bitwise_xor(data, mask).tostring() b = numpy.bitwise_xor(data, mask).tostring()
if ret['length'] % 4: if ret['length'] % 4:
print "Partial unmask" print("Partial unmask")
mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'), mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
offset=header_len, count=(ret['length'] % 4)) offset=header_len, count=(ret['length'] % 4))
data = numpy.frombuffer(buf, dtype=numpy.dtype('B'), data = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
@ -247,20 +279,20 @@ Sec-WebSocket-Accept: %s\r
c = numpy.bitwise_xor(data, mask).tostring() c = numpy.bitwise_xor(data, mask).tostring()
ret['payload'] = b + c ret['payload'] = b + c
else: else:
print "Unmasked frame:", repr(buf) print("Unmasked frame: %s" % repr(buf))
ret['payload'] = buf[(header_len + has_mask * 4):full_len] ret['payload'] = buf[(header_len + has_mask * 4):full_len]
if base64 and ret['opcode'] in [1, 2]: if base64 and ret['opcode'] in [1, 2]:
try: try:
ret['payload'] = b64decode(ret['payload']) ret['payload'] = b64decode(ret['payload'])
except: except:
print "Exception while b64decoding buffer:", repr(buf) print("Exception while b64decoding buffer: %s" %
repr(buf))
raise raise
if ret['opcode'] == 0x08: if ret['opcode'] == 0x08:
if ret['length'] >= 2: if ret['length'] >= 2:
ret['close_code'] = struct.unpack_from( ret['close_code'] = struct.unpack_from(">H", ret['payload'])
">H", ret['payload'])
if ret['length'] > 3: if ret['length'] > 3:
ret['close_reason'] = ret['payload'][2:] ret['close_reason'] = ret['payload'][2:]
@ -268,11 +300,11 @@ Sec-WebSocket-Accept: %s\r
@staticmethod @staticmethod
def encode_hixie(buf): def encode_hixie(buf):
return "\x00" + b64encode(buf) + "\xff" return s2b("\x00" + b2s(b64encode(buf)) + "\xff")
@staticmethod @staticmethod
def decode_hixie(buf): def decode_hixie(buf):
end = buf.find('\xff') end = buf.find(s2b('\xff'))
return {'payload': b64decode(buf[1:end]), return {'payload': b64decode(buf[1:end]),
'left': len(buf) - (end + 1)} 'left': len(buf) - (end + 1)}
@ -288,7 +320,8 @@ Sec-WebSocket-Accept: %s\r
num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1 num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2 num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
return md5(struct.pack('>II8s', num1, num2, key3)).digest() return b2s(md5(struct.pack('>II8s',
int(num1), int(num2), key3)).digest())
# #
# WebSocketServer logging/output functions # WebSocketServer logging/output functions
@ -303,7 +336,7 @@ Sec-WebSocket-Accept: %s\r
def msg(self, msg): def msg(self, msg):
""" Output message with handler_id prefix. """ """ Output message with handler_id prefix. """
if not self.daemon: if not self.daemon:
print "% 3d: %s" % (self.handler_id, msg) print("% 3d: %s" % (self.handler_id, msg))
def vmsg(self, msg): def vmsg(self, msg):
""" Same as msg() but only if verbose. """ """ Same as msg() but only if verbose. """
@ -371,7 +404,7 @@ Sec-WebSocket-Accept: %s\r
if self.version.startswith("hybi"): if self.version.startswith("hybi"):
frame = self.decode_hybi(buf, base64=self.base64) frame = self.decode_hybi(buf, base64=self.base64)
#print "Received buf: %s, frame: %s" % (repr(buf), frame) #print("Received buf: %s, frame: %s" % (repr(buf), frame))
if frame['payload'] == None: if frame['payload'] == None:
# Incomplete/partial frame # Incomplete/partial frame
@ -395,7 +428,7 @@ Sec-WebSocket-Accept: %s\r
buf = buf[2:] buf = buf[2:]
continue # No-op continue # No-op
elif buf.count('\xff') == 0: elif buf.count(s2b('\xff')) == 0:
# Partial frame # Partial frame
self.traffic("}.") self.traffic("}.")
self.recv_part = buf self.recv_part = buf
@ -418,7 +451,7 @@ Sec-WebSocket-Accept: %s\r
""" Send a WebSocket orderly close frame. """ """ Send a WebSocket orderly close frame. """
if self.version.startswith("hybi"): if self.version.startswith("hybi"):
msg = '' msg = s2b('')
if code != None: if code != None:
msg = struct.pack(">H%ds" % (len(reason)), code) msg = struct.pack(">H%ds" % (len(reason)), code)
@ -426,7 +459,7 @@ Sec-WebSocket-Accept: %s\r
self.client.send(buf) self.client.send(buf)
elif self.version == "hixie-76": elif self.version == "hixie-76":
buf = self.encode_hixie('\xff\x00') buf = s2b('\xff\x00')
self.client.send(buf) self.client.send(buf)
# No orderly close for 75 # No orderly close for 75
@ -462,14 +495,16 @@ Sec-WebSocket-Accept: %s\r
if handshake == "": if handshake == "":
raise self.EClose("ignoring empty handshake") raise self.EClose("ignoring empty handshake")
elif handshake.startswith("<policy-file-request/>"): elif handshake.startswith(s2b("<policy-file-request/>")):
# Answer Flash policy request # Answer Flash policy request
handshake = sock.recv(1024) handshake = sock.recv(1024)
sock.send(self.policy_response) sock.send(s2b(self.policy_response))
raise self.EClose("Sending flash policy response") raise self.EClose("Sending flash policy response")
elif handshake[0] in ("\x16", "\x80"): elif handshake[0] in ("\x16", "\x80"):
# SSL wrap the connection # SSL wrap the connection
if not ssl:
raise self.EClose("SSL connection but no 'ssl' module")
if not os.path.exists(self.cert): if not os.path.exists(self.cert):
raise self.EClose("SSL connection but '%s' not found" raise self.EClose("SSL connection but '%s' not found"
% self.cert) % self.cert)
@ -479,7 +514,8 @@ Sec-WebSocket-Accept: %s\r
server_side=True, server_side=True,
certfile=self.cert, certfile=self.cert,
keyfile=self.key) keyfile=self.key)
except ssl.SSLError, x: except ssl.SSLError:
_, x, _ = sys.exc_info()
if x.args[0] == ssl.SSL_ERROR_EOF: if x.args[0] == ssl.SSL_ERROR_EOF:
raise self.EClose("") raise self.EClose("")
else: else:
@ -519,8 +555,8 @@ Sec-WebSocket-Accept: %s\r
if ver: if ver:
# HyBi/IETF version of the protocol # HyBi/IETF version of the protocol
if not numpy or not ctypes: if sys.hexversion < 0x2060000 or not numpy:
self.EClose("Python numpy and ctypes modules required for HyBi-07 or greater") raise self.EClose("Python >= 2.6 and numpy module is required for HyBi-07 or greater")
if ver == '7': if ver == '7':
self.version = "hybi-07" self.version = "hybi-07"
@ -538,7 +574,7 @@ Sec-WebSocket-Accept: %s\r
raise self.EClose("Client must support 'binary' or 'base64' protocol") raise self.EClose("Client must support 'binary' or 'base64' protocol")
# Generate the hash value for the accept header # Generate the hash value for the accept header
accept = b64encode(sha1(key + self.GUID).digest()) accept = b64encode(sha1(s2b(key + self.GUID)).digest())
response = self.server_handshake_hybi % accept response = self.server_handshake_hybi % accept
if self.base64: if self.base64:
@ -577,7 +613,7 @@ Sec-WebSocket-Accept: %s\r
# Send server WebSockets handshake response # Send server WebSockets handshake response
#self.msg("sending response [%s]" % response) #self.msg("sending response [%s]" % response)
retsock.send(response) retsock.send(s2b(response))
# Return the WebSockets socket which may be SSL wrapped # Return the WebSockets socket which may be SSL wrapped
return retsock return retsock
@ -595,9 +631,8 @@ Sec-WebSocket-Accept: %s\r
#self.vmsg("Running poll()") #self.vmsg("Running poll()")
pass pass
def top_SIGCHLD(self, sig, stack): def fallback_SIGCHLD(self, sig, stack):
# Reap zombies after calling child SIGCHLD handler # Reap zombies when using os.fork() (python 2.4)
self.do_SIGCHLD(sig, stack)
self.vmsg("Got SIGCHLD, reaping zombies") self.vmsg("Got SIGCHLD, reaping zombies")
try: try:
result = os.waitpid(-1, os.WNOHANG) result = os.waitpid(-1, os.WNOHANG)
@ -607,14 +642,37 @@ Sec-WebSocket-Accept: %s\r
except (OSError): except (OSError):
pass pass
def do_SIGCHLD(self, sig, stack):
pass
def do_SIGINT(self, sig, stack): def do_SIGINT(self, sig, stack):
self.msg("Got SIGINT, exiting") self.msg("Got SIGINT, exiting")
sys.exit(0) sys.exit(0)
def new_client(self, client): def top_new_client(self, startsock, address):
""" Do something with a WebSockets client connection. """
# Initialize per client settings
self.send_parts = []
self.recv_part = None
self.base64 = False
# handler process
try:
try:
self.client = self.do_handshake(startsock, address)
self.new_client()
except self.EClose:
_, exc, _ = sys.exc_info()
# Connection was not a WebSockets connection
if exc.args[0]:
self.msg("%s: %s" % (address[0], exc.args[0]))
except Exception:
_, exc, _ = sys.exc_info()
self.msg("handler exception: %s" % str(exc))
if self.verbose:
self.msg(traceback.format_exc())
finally:
if self.client and self.client != startsock:
self.client.close()
def new_client(self):
""" Do something with a WebSockets client connection. """ """ Do something with a WebSockets client connection. """
raise("WebSocketServer.new_client() must be overloaded") raise("WebSocketServer.new_client() must be overloaded")
@ -636,9 +694,11 @@ Sec-WebSocket-Accept: %s\r
self.started() # Some things need to happen after daemonizing self.started() # Some things need to happen after daemonizing
# Reep zombies # Allow override of SIGINT
signal.signal(signal.SIGCHLD, self.top_SIGCHLD)
signal.signal(signal.SIGINT, self.do_SIGINT) signal.signal(signal.SIGINT, self.do_SIGINT)
if not Process:
# os.fork() (python 2.4) child reaper
signal.signal(signal.SIGCHLD, self.fallback_SIGCHLD)
while True: while True:
try: try:
@ -655,9 +715,12 @@ Sec-WebSocket-Accept: %s\r
startsock, address = lsock.accept() startsock, address = lsock.accept()
else: else:
continue continue
except Exception, exc: except Exception:
_, exc, _ = sys.exc_info()
if hasattr(exc, 'errno'): if hasattr(exc, 'errno'):
err = exc.errno err = exc.errno
elif hasattr(exc, 'args'):
err = exc.args[0]
else: else:
err = exc[0] err = exc[0]
if err == errno.EINTR: if err == errno.EINTR:
@ -666,41 +729,42 @@ Sec-WebSocket-Accept: %s\r
else: else:
raise raise
self.vmsg('%s: forking handler' % address[0]) if Process:
pid = os.fork() self.vmsg('%s: new handler Process' % address[0])
p = Process(target=self.top_new_client,
if pid == 0: args=(startsock, address))
# Initialize per client settings p.start()
self.send_parts = [] # child will not return
self.recv_part = None
self.base64 = False
# handler process
self.client = self.do_handshake(
startsock, address)
self.new_client()
else: else:
# parent process # python 2.4
self.handler_id += 1 self.vmsg('%s: forking handler' % address[0])
pid = os.fork()
if pid == 0:
# child handler process
self.top_new_client(startsock, address)
break # child process exits
except self.EClose, exc: # parent process
# Connection was not a WebSockets connection self.handler_id += 1
if exc.args[0]:
self.msg("%s: %s" % (address[0], exc.args[0])) except KeyboardInterrupt:
except KeyboardInterrupt, exc: _, exc, _ = sys.exc_info()
print("In KeyboardInterrupt")
pass pass
except Exception, exc: except SystemExit:
_, exc, _ = sys.exc_info()
print("In SystemExit")
break
except Exception:
_, exc, _ = sys.exc_info()
self.msg("handler exception: %s" % str(exc)) self.msg("handler exception: %s" % str(exc))
if self.verbose: if self.verbose:
self.msg(traceback.format_exc()) self.msg(traceback.format_exc())
finally: finally:
if self.client and self.client != startsock:
self.client.close()
if startsock: if startsock:
startsock.close() startsock.close()
if pid == 0:
break # Child process exits
# HTTP handler with WebSocket upgrade support # HTTP handler with WebSocket upgrade support
class WSRequestHandler(SimpleHTTPRequestHandler): class WSRequestHandler(SimpleHTTPRequestHandler):
@ -709,13 +773,13 @@ class WSRequestHandler(SimpleHTTPRequestHandler):
SimpleHTTPRequestHandler.__init__(self, req, addr, object()) SimpleHTTPRequestHandler.__init__(self, req, addr, object())
def do_GET(self): def do_GET(self):
if (self.headers.has_key('upgrade') and if (self.headers.get('upgrade') and
self.headers.get('upgrade').lower() == 'websocket'): self.headers.get('upgrade').lower() == 'websocket'):
if (self.headers.get('sec-websocket-key1') or if (self.headers.get('sec-websocket-key1') or
self.headers.get('websocket-key1')): self.headers.get('websocket-key1')):
# For Hixie-76 read out the key hash # For Hixie-76 read out the key hash
self.headers.dict['key3'] = self.rfile.read(8) self.headers.__setitem__('key3', self.rfile.read(8))
# Just indicate that an WebSocket upgrade is needed # Just indicate that an WebSocket upgrade is needed
self.last_code = 101 self.last_code = 101

16
websockify
View File

@ -2,7 +2,7 @@
''' '''
A WebSocket to TCP socket proxy with support for "wss://" encryption. A WebSocket to TCP socket proxy with support for "wss://" encryption.
Copyright 2010 Joel Martin Copyright 2011 Joel Martin
Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3) Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
You can make a cert/key with openssl using: You can make a cert/key with openssl using:
@ -74,7 +74,7 @@ Traffic Legend:
WebSocketServer.__init__(self, *args, **kwargs) WebSocketServer.__init__(self, *args, **kwargs)
def run_wrap_cmd(self): def run_wrap_cmd(self):
print "Starting '%s'" % " ".join(self.wrap_cmd) print("Starting '%s'" % " ".join(self.wrap_cmd))
self.wrap_times.append(time.time()) self.wrap_times.append(time.time())
self.wrap_times.pop(0) self.wrap_times.pop(0)
self.cmd = subprocess.Popen( self.cmd = subprocess.Popen(
@ -88,14 +88,14 @@ Traffic Legend:
# Need to call wrapped command after daemonization so we can # Need to call wrapped command after daemonization so we can
# know when the wrapped command exits # know when the wrapped command exits
if self.wrap_cmd: if self.wrap_cmd:
print " - proxying from %s:%s to '%s' (port %s)\n" % ( print(" - proxying from %s:%s to '%s' (port %s)\n" % (
self.listen_host, self.listen_port, self.listen_host, self.listen_port,
" ".join(self.wrap_cmd), self.target_port) " ".join(self.wrap_cmd), self.target_port))
self.run_wrap_cmd() self.run_wrap_cmd()
else: else:
print " - proxying from %s:%s to %s:%s\n" % ( print(" - proxying from %s:%s to %s:%s\n" % (
self.listen_host, self.listen_port, self.listen_host, self.listen_port,
self.target_host, self.target_port) self.target_host, self.target_port))
def poll(self): def poll(self):
# If we are wrapping a command, check it's status # If we are wrapping a command, check it's status
@ -118,7 +118,7 @@ Traffic Legend:
if (now - avg) < 10: if (now - avg) < 10:
# 3 times in the last 10 seconds # 3 times in the last 10 seconds
if self.spawn_message: if self.spawn_message:
print "Command respawning too fast" print("Command respawning too fast")
self.spawn_message = False self.spawn_message = False
else: else:
self.run_wrap_cmd() self.run_wrap_cmd()
@ -154,7 +154,7 @@ Traffic Legend:
tsock.connect((self.target_host, self.target_port)) tsock.connect((self.target_host, self.target_port))
if self.verbose and not self.daemon: if self.verbose and not self.daemon:
print self.traffic_legend print(self.traffic_legend)
# Start proxying # Start proxying
try: try: