utils/websocket.py

   1 #!/usr/bin/python
   2
   3 '''
   4 Python WebSocket library with support for "wss://" encryption.
   5
   6 You can make a cert/key with openssl using:
   7 openssl req -new -x509 -days 365 -nodes -out self.pem -keyout self.pem
   8 as taken from http://docs.python.org/dev/library/ssl.html#certificates
   9
  10 '''
  11
  12 import sys, socket, ssl, struct, traceback
  13 import os, resource, errno, signal # daemonizing
  14 from base64 import b64encode, b64decode
  15 from hashlib import md5
  16
  17 settings = {
  18     'listen_host' : '',
  19     'listen_port' : None,
  20     'handler'     : None,
  21     'cert'        : None,
  22     'ssl_only'    : False,
  23     'daemon'      : True,
  24     'record'      : None, }
  25 client_settings = {
  26     'b64encode'   : False,
  27     'seq_num'     : False, }
  28
  29 send_seq = 0
  30
  31 server_handshake = """HTTP/1.1 101 Web Socket Protocol Handshake\r
  32 Upgrade: WebSocket\r
  33 Connection: Upgrade\r
  34 %sWebSocket-Origin: %s\r
  35 %sWebSocket-Location: %s://%s%s\r
  36 %sWebSocket-Protocol: sample\r
  37 \r
  38 %s"""
  39
  40 policy_response = """<cross-domain-policy><allow-access-from domain="*" to-ports="*" /></cross-domain-policy>\n"""
  41
  42 def traffic(token="."):
  43     sys.stdout.write(token)
  44     sys.stdout.flush()
  45
  46 def decode(buf):
  47     """ Parse out WebSocket packets. """
  48     if buf.count('\xff') > 1:
  49         if client_settings['b64encode']:
  50             return [b64decode(d[1:]) for d in buf.split('\xff')]
  51         else:
  52             # Modified UTF-8 decode
  53             return [d[1:].replace("\xc4\x80", "\x00").decode('utf-8').encode('latin-1') for d in buf.split('\xff')]
  54     else:
  55         if client_settings['b64encode']:
  56             return [b64decode(buf[1:-1])]
  57         else:
  58             return [buf[1:-1].replace("\xc4\x80", "\x00").decode('utf-8').encode('latin-1')]
  59
  60 def encode(buf):
  61     global send_seq
  62     if client_settings['b64encode']:
  63         buf = b64encode(buf)
  64     else:
  65         # Modified UTF-8 encode
  66         buf = buf.decode('latin-1').encode('utf-8').replace("\x00", "\xc4\x80")
  67
  68     if client_settings['seq_num']:
  69         send_seq += 1
  70         return "\x00%d:%s\xff" % (send_seq-1, buf)
  71     else:
  72         return "\x00%s\xff" % buf
  73
  74
  75 def do_handshake(sock):
  76     global client_settings, send_seq
  77
  78     client_settings['b64encode'] = False
  79     client_settings['seq_num'] = False
  80     send_seq = 0
  81
  82     # Peek, but don't read the data
  83     handshake = sock.recv(1024, socket.MSG_PEEK)
  84     #print "Handshake [%s]" % repr(handshake)
  85     if handshake.startswith("<policy-file-request/>"):
  86         handshake = sock.recv(1024)
  87         print "Sending flash policy response"
  88         sock.send(policy_response)
  89         sock.close()
  90         return False
  91     elif handshake.startswith("\x16"):
  92         retsock = ssl.wrap_socket(
  93                 sock,
  94                 server_side=True,
  95                 certfile=settings['cert'],
  96                 ssl_version=ssl.PROTOCOL_TLSv1)
  97         scheme = "wss"
  98         print "  using SSL/TLS"
  99     elif settings['ssl_only']:
 100         print "Non-SSL connection disallowed"
 101         sock.close()
 102         return False
 103     else:
 104         retsock = sock
 105         scheme = "ws"
 106         print "  using plain (not SSL) socket"
 107     handshake = retsock.recv(4096)
 108     #print "handshake: " + repr(handshake)
 109     h = parse_handshake(handshake)
 110
 111     # Parse client settings from the GET path
 112     cvars = h['path'].partition('?')[2].partition('#')[0].split('&')
 113     for cvar in [c for c in cvars if c]:
 114         name, _, val = cvar.partition('=')
 115         if name not in ['b64encode', 'seq_num']: continue
 116         value = val and val or True
 117         client_settings[name] = value
 118         print "  %s=%s" % (name, value)
 119
 120     if h.get('key3'):
 121         trailer = gen_md5(h)
 122         pre = "Sec-"
 123     else:
 124         trailer = ""
 125         pre = ""
 126
 127     response = server_handshake % (pre, h['Origin'], pre, scheme,
 128             h['Host'], h['path'], pre, trailer)
 129
 130     #print "sending response:", repr(response)
 131     retsock.send(response)
 132     return retsock
 133
 134 def parse_handshake(handshake):
 135     ret = {}
 136     req_lines = handshake.split("\r\n")
 137     if not req_lines[0].startswith("GET "):
 138         raise "Invalid handshake: no GET request line"
 139     ret['path'] = req_lines[0].split(" ")[1]
 140     for line in req_lines[1:]:
 141         if line == "": break
 142         var, delim, val = line.partition(": ")
 143         ret[var] = val
 144
 145     if req_lines[-2] == "":
 146         ret['key3'] = req_lines[-1]
 147
 148     return ret
 149
 150 def gen_md5(keys):
 151     key1 = keys['Sec-WebSocket-Key1']
 152     key2 = keys['Sec-WebSocket-Key2']
 153     key3 = keys['key3']
 154     spaces1 = key1.count(" ")
 155     spaces2 = key2.count(" ")
 156     num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
 157     num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
 158
 159     return md5(struct.pack('>II8s', num1, num2, key3)).digest()
 160
 161 def daemonize():
 162     os.umask(0)
 163     os.chdir('/')
 164     os.setgid(os.getgid())  # relinquish elevations
 165     os.setuid(os.getuid())  # relinquish elevations
 166
 167     # Double fork to daemonize
 168     if os.fork() > 0: os._exit(0)  # Parent exits
 169     os.setsid()                    # Obtain new process group
 170     if os.fork() > 0: os._exit(0)  # Parent exits
 171
 172     # Signal handling
 173     def terminate(a,b): os._exit(0)
 174     signal.signal(signal.SIGTERM, terminate)
 175     signal.signal(signal.SIGINT, signal.SIG_IGN)
 176
 177     # Close open files
 178     maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
 179     if maxfd == resource.RLIM_INFINITY: maxfd = 256
 180     for fd in reversed(range(maxfd)):
 181         try:
 182             os.close(fd)
 183         except OSError, exc:
 184             if exc.errno != errno.EBADF: raise
 185
 186     # Redirect I/O to /dev/null
 187     os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdin.fileno())
 188     os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdout.fileno())
 189     os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno())
 190
 191
 192 def start_server():
 193
 194     if settings['daemon']: daemonize()
 195
 196     lsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 197     lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 198     lsock.bind((settings['listen_host'], settings['listen_port']))
 199     lsock.listen(100)
 200     while True:
 201         try:
 202             csock = startsock = None
 203             print 'waiting for connection on port %s' % settings['listen_port']
 204             startsock, address = lsock.accept()
 205             print 'Got client connection from %s' % address[0]
 206             csock = do_handshake(startsock)
 207             if not csock: continue
 208
 209             settings['handler'](csock)
 210
 211         except Exception:
 212             print "Ignoring exception:"
 213             print traceback.format_exc()
 214             if csock: csock.close()
 215             if startsock and startsock != csock: startsock.close()