utils/websocket.py

   1 #!/usr/bin/env python
   2
   3 '''
   4 Python WebSocket library with support for "wss://" encryption.
   5 Copyright 2011 Joel Martin
   6 Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
   7
   8 Supports following protocol versions:
   9     - http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-75
  10     - http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
  11     - http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-10
  12
  13 You can make a cert/key with openssl using:
  14 openssl req -new -x509 -days 365 -nodes -out self.pem -keyout self.pem
  15 as taken from http://docs.python.org/dev/library/ssl.html#certificates
  16
  17 '''
  18
  19 import os, sys, time, errno, signal, socket, traceback, select
  20 import array, struct
  21 from cgi import parse_qsl
  22 from base64 import b64encode, b64decode
  23
  24 # Imports that vary by python version
  25
  26 # python 3.0 differences
  27 if sys.hexversion > 0x3000000:
  28     b2s = lambda buf: buf.decode('latin_1')
  29     s2b = lambda s: s.encode('latin_1')
  30     s2a = lambda s: s
  31 else:
  32     b2s = lambda buf: buf  # No-op
  33     s2b = lambda s: s      # No-op
  34     s2a = lambda s: [ord(c) for c in s]
  35 try:    from io import StringIO
  36 except: from cStringIO import StringIO
  37 try:    from http.server import SimpleHTTPRequestHandler
  38 except: from SimpleHTTPServer import SimpleHTTPRequestHandler
  39 try:    from urllib.parse import urlsplit
  40 except: from urlparse import urlsplit
  41
  42 # python 2.6 differences
  43 try:    from hashlib import md5, sha1
  44 except: from md5 import md5; from sha import sha as sha1
  45
  46 # python 2.5 differences
  47 try:
  48     from struct import pack, unpack_from
  49 except:
  50     from struct import pack
  51     def unpack_from(fmt, buf, offset=0):
  52         slice = buffer(buf, offset, struct.calcsize(fmt))
  53         return struct.unpack(fmt, slice)
  54
  55 # Degraded functionality if these imports are missing
  56 for mod, sup in [('numpy', 'HyBi protocol'), ('ssl', 'TLS/SSL/wss'),
  57         ('multiprocessing', 'Multi-Processing'),
  58         ('resource', 'daemonizing')]:
  59     try:
  60         globals()[mod] = __import__(mod)
  61     except ImportError:
  62         globals()[mod] = None
  63         print("WARNING: no '%s' module, %s is slower or disabled" % (
  64             mod, sup))
  65 if multiprocessing and sys.platform == 'win32':
  66     # make sockets pickle-able/inheritable
  67     import multiprocessing.reduction
  68
  69
  70 class WebSocketServer(object):
  71     """
  72     WebSockets server class.
  73     Must be sub-classed with new_client method definition.
  74     """
  75
  76     buffer_size = 65536
  77
  78     server_handshake_hixie = """HTTP/1.1 101 Web Socket Protocol Handshake\r
  79 Upgrade: WebSocket\r
  80 Connection: Upgrade\r
  81 %sWebSocket-Origin: %s\r
  82 %sWebSocket-Location: %s://%s%s\r
  83 """
  84
  85     server_handshake_hybi = """HTTP/1.1 101 Switching Protocols\r
  86 Upgrade: websocket\r
  87 Connection: Upgrade\r
  88 Sec-WebSocket-Accept: %s\r
  89 """
  90
  91     GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
  92
  93     policy_response = """<cross-domain-policy><allow-access-from domain="*" to-ports="*" /></cross-domain-policy>\n"""
  94
  95     # An exception before the WebSocket connection was established
  96     class EClose(Exception):
  97         pass
  98
  99     # An exception while the WebSocket client was connected
 100     class CClose(Exception):
 101         pass
 102
 103     def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
 104             verbose=False, cert='', key='', ssl_only=None,
 105             daemon=False, record='', web='',
 106             run_once=False, timeout=0):
 107
 108         # settings
 109         self.verbose        = verbose
 110         self.listen_host    = listen_host
 111         self.listen_port    = listen_port
 112         self.ssl_only       = ssl_only
 113         self.daemon         = daemon
 114         self.run_once       = run_once
 115         self.timeout        = timeout
 116
 117         self.launch_time    = time.time()
 118         self.ws_connection  = False
 119         self.handler_id     = 1
 120
 121         # Make paths settings absolute
 122         self.cert = os.path.abspath(cert)
 123         self.key = self.web = self.record = ''
 124         if key:
 125             self.key = os.path.abspath(key)
 126         if web:
 127             self.web = os.path.abspath(web)
 128         if record:
 129             self.record = os.path.abspath(record)
 130
 131         if self.web:
 132             os.chdir(self.web)
 133
 134         # Sanity checks
 135         if not ssl and self.ssl_only:
 136             raise Exception("No 'ssl' module and SSL-only specified")
 137         if self.daemon and not resource:
 138             raise Exception("Module 'resource' required to daemonize")
 139
 140         # Show configuration
 141         print("WebSocket server settings:")
 142         print("  - Listen on %s:%s" % (
 143                 self.listen_host, self.listen_port))
 144         print("  - Flash security policy server")
 145         if self.web:
 146             print("  - Web server. Web root: %s" % self.web)
 147         if ssl:
 148             if os.path.exists(self.cert):
 149                 print("  - SSL/TLS support")
 150                 if self.ssl_only:
 151                     print("  - Deny non-SSL/TLS connections")
 152             else:
 153                 print("  - No SSL/TLS support (no cert file)")
 154         else:
 155             print("  - No SSL/TLS support (no 'ssl' module)")
 156         if self.daemon:
 157             print("  - Backgrounding (daemon)")
 158         if self.record:
 159             print("  - Recording to '%s.*'" % self.record)
 160
 161     #
 162     # WebSocketServer static methods
 163     #
 164
 165     @staticmethod
 166     def socket(host, port=None, connect=False, prefer_ipv6=False):
 167         """ Resolve a host (and optional port) to an IPv4 or IPv6
 168         address. Create a socket. Bind to it if listen is set,
 169         otherwise connect to it. Return the socket.
 170         """
 171         flags = 0
 172         if host == '':
 173             host = None
 174         if connect and not port:
 175             raise Exception("Connect mode requires a port")
 176         if not connect:
 177             flags = flags | socket.AI_PASSIVE
 178         addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
 179                 socket.IPPROTO_TCP, flags)
 180         if not addrs:
 181             raise Exception("Could resolve host '%s'" % host)
 182         addrs.sort(key=lambda x: x[0])
 183         if prefer_ipv6:
 184             addrs.reverse()
 185         sock = socket.socket(addrs[0][0], addrs[0][1])
 186         if connect:
 187             sock.connect(addrs[0][4])
 188         else:
 189             sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 190             sock.bind(addrs[0][4])
 191             sock.listen(100)
 192         return sock
 193
 194     @staticmethod
 195     def daemonize(keepfd=None, chdir='/'):
 196         os.umask(0)
 197         if chdir:
 198             os.chdir(chdir)
 199         else:
 200             os.chdir('/')
 201         os.setgid(os.getgid())  # relinquish elevations
 202         os.setuid(os.getuid())  # relinquish elevations
 203
 204         # Double fork to daemonize
 205         if os.fork() > 0: os._exit(0)  # Parent exits
 206         os.setsid()                    # Obtain new process group
 207         if os.fork() > 0: os._exit(0)  # Parent exits
 208
 209         # Signal handling
 210         def terminate(a,b): os._exit(0)
 211         signal.signal(signal.SIGTERM, terminate)
 212         signal.signal(signal.SIGINT, signal.SIG_IGN)
 213
 214         # Close open files
 215         maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
 216         if maxfd == resource.RLIM_INFINITY: maxfd = 256
 217         for fd in reversed(range(maxfd)):
 218             try:
 219                 if fd != keepfd:
 220                     os.close(fd)
 221             except OSError:
 222                 _, exc, _ = sys.exc_info()
 223                 if exc.errno != errno.EBADF: raise
 224
 225         # Redirect I/O to /dev/null
 226         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdin.fileno())
 227         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdout.fileno())
 228         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno())
 229
 230     @staticmethod
 231     def unmask(buf, f):
 232         pstart = f['hlen'] + 4
 233         pend = pstart + f['length']
 234         if numpy:
 235             b = c = s2b('')
 236             if f['length'] >= 4:
 237                 mask = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
 238                         offset=f['hlen'], count=1)
 239                 data = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
 240                         offset=pstart, count=int(f['length'] / 4))
 241                 #b = numpy.bitwise_xor(data, mask).data
 242                 b = numpy.bitwise_xor(data, mask).tostring()
 243
 244             if f['length'] % 4:
 245                 #print("Partial unmask")
 246                 mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
 247                         offset=f['hlen'], count=(f['length'] % 4))
 248                 data = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
 249                         offset=pend - (f['length'] % 4),
 250                         count=(f['length'] % 4))
 251                 c = numpy.bitwise_xor(data, mask).tostring()
 252             return b + c
 253         else:
 254             # Slower fallback
 255             data = array.array('B')
 256             mask = s2a(f['mask'])
 257             data.fromstring(buf[pstart:pend])
 258             for i in range(len(data)):
 259                 data[i] ^= mask[i % 4]
 260             return data.tostring()
 261
 262     @staticmethod
 263     def encode_hybi(buf, opcode, base64=False):
 264         """ Encode a HyBi style WebSocket frame.
 265         Optional opcode:
 266             0x0 - continuation
 267             0x1 - text frame (base64 encode buf)
 268             0x2 - binary frame (use raw buf)
 269             0x8 - connection close
 270             0x9 - ping
 271             0xA - pong
 272         """
 273         if base64:
 274             buf = b64encode(buf)
 275
 276         b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
 277         payload_len = len(buf)
 278         if payload_len <= 125:
 279             header = pack('>BB', b1, payload_len)
 280         elif payload_len > 125 and payload_len < 65536:
 281             header = pack('>BBH', b1, 126, payload_len)
 282         elif payload_len >= 65536:
 283             header = pack('>BBQ', b1, 127, payload_len)
 284
 285         #print("Encoded: %s" % repr(header + buf))
 286
 287         return header + buf, len(header), 0
 288
 289     @staticmethod
 290     def decode_hybi(buf, base64=False):
 291         """ Decode HyBi style WebSocket packets.
 292         Returns:
 293             {'fin'          : 0_or_1,
 294              'opcode'       : number,
 295              'mask'         : 32_bit_number,
 296              'hlen'         : header_bytes_number,
 297              'length'       : payload_bytes_number,
 298              'payload'      : decoded_buffer,
 299              'left'         : bytes_left_number,
 300              'close_code'   : number,
 301              'close_reason' : string}
 302         """
 303
 304         f = {'fin'          : 0,
 305              'opcode'       : 0,
 306              'mask'         : 0,
 307              'hlen'         : 2,
 308              'length'       : 0,
 309              'payload'      : None,
 310              'left'         : 0,
 311              'close_code'   : 1000,
 312              'close_reason' : ''}
 313
 314         blen = len(buf)
 315         f['left'] = blen
 316
 317         if blen < f['hlen']:
 318             return f # Incomplete frame header
 319
 320         b1, b2 = unpack_from(">BB", buf)
 321         f['opcode'] = b1 & 0x0f
 322         f['fin'] = (b1 & 0x80) >> 7
 323         has_mask = (b2 & 0x80) >> 7
 324
 325         f['length'] = b2 & 0x7f
 326
 327         if f['length'] == 126:
 328             f['hlen'] = 4
 329             if blen < f['hlen']:
 330                 return f # Incomplete frame header
 331             (f['length'],) = unpack_from('>xxH', buf)
 332         elif f['length'] == 127:
 333             f['hlen'] = 10
 334             if blen < f['hlen']:
 335                 return f # Incomplete frame header
 336             (f['length'],) = unpack_from('>xxQ', buf)
 337
 338         full_len = f['hlen'] + has_mask * 4 + f['length']
 339
 340         if blen < full_len: # Incomplete frame
 341             return f # Incomplete frame header
 342
 343         # Number of bytes that are part of the next frame(s)
 344         f['left'] = blen - full_len
 345
 346         # Process 1 frame
 347         if has_mask:
 348             # unmask payload
 349             f['mask'] = buf[f['hlen']:f['hlen']+4]
 350             f['payload'] = WebSocketServer.unmask(buf, f)
 351         else:
 352             print("Unmasked frame: %s" % repr(buf))
 353             f['payload'] = buf[(f['hlen'] + has_mask * 4):full_len]
 354
 355         if base64 and f['opcode'] in [1, 2]:
 356             try:
 357                 f['payload'] = b64decode(f['payload'])
 358             except:
 359                 print("Exception while b64decoding buffer: %s" %
 360                         repr(buf))
 361                 raise
 362
 363         if f['opcode'] == 0x08:
 364             if f['length'] >= 2:
 365                 f['close_code'] = unpack_from(">H", f['payload'])[0]
 366             if f['length'] > 3:
 367                 f['close_reason'] = f['payload'][2:]
 368
 369         return f
 370
 371     @staticmethod
 372     def encode_hixie(buf):
 373         return s2b("\x00" + b2s(b64encode(buf)) + "\xff"), 1, 1
 374
 375     @staticmethod
 376     def decode_hixie(buf):
 377         end = buf.find(s2b('\xff'))
 378         return {'payload': b64decode(buf[1:end]),
 379                 'hlen': 1,
 380                 'length': end - 1,
 381                 'left': len(buf) - (end + 1)}
 382
 383
 384     @staticmethod
 385     def gen_md5(keys):
 386         """ Generate hash value for WebSockets hixie-76. """
 387         key1 = keys['Sec-WebSocket-Key1']
 388         key2 = keys['Sec-WebSocket-Key2']
 389         key3 = keys['key3']
 390         spaces1 = key1.count(" ")
 391         spaces2 = key2.count(" ")
 392         num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
 393         num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
 394
 395         return b2s(md5(pack('>II8s',
 396             int(num1), int(num2), key3)).digest())
 397
 398     #
 399     # WebSocketServer logging/output functions
 400     #
 401
 402     def traffic(self, token="."):
 403         """ Show traffic flow in verbose mode. """
 404         if self.verbose and not self.daemon:
 405             sys.stdout.write(token)
 406             sys.stdout.flush()
 407
 408     def msg(self, msg):
 409         """ Output message with handler_id prefix. """
 410         if not self.daemon:
 411             print("% 3d: %s" % (self.handler_id, msg))
 412
 413     def vmsg(self, msg):
 414         """ Same as msg() but only if verbose. """
 415         if self.verbose:
 416             self.msg(msg)
 417
 418     #
 419     # Main WebSocketServer methods
 420     #
 421     def send_frames(self, bufs=None):
 422         """ Encode and send WebSocket frames. Any frames already
 423         queued will be sent first. If buf is not set then only queued
 424         frames will be sent. Returns the number of pending frames that
 425         could not be fully sent. If returned pending frames is greater
 426         than 0, then the caller should call again when the socket is
 427         ready. """
 428
 429         tdelta = int(time.time()*1000) - self.start_time
 430
 431         if bufs:
 432             for buf in bufs:
 433                 if self.version.startswith("hybi"):
 434                     if self.base64:
 435                         encbuf, lenhead, lentail = self.encode_hybi(
 436                                 buf, opcode=1, base64=True)
 437                     else:
 438                         encbuf, lenhead, lentail = self.encode_hybi(
 439                                 buf, opcode=2, base64=False)
 440
 441                 else:
 442                     encbuf, lenhead, lentail = self.encode_hixie(buf)
 443
 444                 if self.rec:
 445                     self.rec.write("%s,\n" %
 446                             repr("{%s{" % tdelta
 447                                 + encbuf[lenhead:-lentail]))
 448
 449                 self.send_parts.append(encbuf)
 450
 451         while self.send_parts:
 452             # Send pending frames
 453             buf = self.send_parts.pop(0)
 454             sent = self.client.send(buf)
 455
 456             if sent == len(buf):
 457                 self.traffic("<")
 458             else:
 459                 self.traffic("<.")
 460                 self.send_parts.insert(0, buf[sent:])
 461                 break
 462
 463         return len(self.send_parts)
 464
 465     def recv_frames(self):
 466         """ Receive and decode WebSocket frames.
 467
 468         Returns:
 469             (bufs_list, closed_string)
 470         """
 471
 472         closed = False
 473         bufs = []
 474         tdelta = int(time.time()*1000) - self.start_time
 475
 476         buf = self.client.recv(self.buffer_size)
 477         if len(buf) == 0:
 478             closed = {'code': 1000, 'reason': "Client closed abruptly"}
 479             return bufs, closed
 480
 481         if self.recv_part:
 482             # Add partially received frames to current read buffer
 483             buf = self.recv_part + buf
 484             self.recv_part = None
 485
 486         while buf:
 487             if self.version.startswith("hybi"):
 488
 489                 frame = self.decode_hybi(buf, base64=self.base64)
 490                 #print("Received buf: %s, frame: %s" % (repr(buf), frame))
 491
 492                 if frame['payload'] == None:
 493                     # Incomplete/partial frame
 494                     self.traffic("}.")
 495                     if frame['left'] > 0:
 496                         self.recv_part = buf[-frame['left']:]
 497                     break
 498                 else:
 499                     if frame['opcode'] == 0x8: # connection close
 500                         closed = {'code': frame['close_code'],
 501                                   'reason': frame['close_reason']}
 502                         break
 503
 504             else:
 505                 if buf[0:2] == s2b('\xff\x00'):
 506                     closed = {'code': 1000,
 507                               'reason': "Client sent orderly close frame"}
 508                     break
 509
 510                 elif buf[0:2] == s2b('\x00\xff'):
 511                     buf = buf[2:]
 512                     continue # No-op
 513
 514                 elif buf.count(s2b('\xff')) == 0:
 515                     # Partial frame
 516                     self.traffic("}.")
 517                     self.recv_part = buf
 518                     break
 519
 520                 frame = self.decode_hixie(buf)
 521
 522             self.traffic("}")
 523
 524             if self.rec:
 525                 start = frame['hlen']
 526                 end = frame['hlen'] + frame['length']
 527                 self.rec.write("%s,\n" %
 528                         repr("}%s}" % tdelta + buf[start:end]))
 529
 530
 531             bufs.append(frame['payload'])
 532
 533             if frame['left']:
 534                 buf = buf[-frame['left']:]
 535             else:
 536                 buf = ''
 537
 538         return bufs, closed
 539
 540     def send_close(self, code=1000, reason=''):
 541         """ Send a WebSocket orderly close frame. """
 542
 543         if self.version.startswith("hybi"):
 544             msg = pack(">H%ds" % len(reason), code, reason)
 545
 546             buf, h, t = self.encode_hybi(msg, opcode=0x08, base64=False)
 547             self.client.send(buf)
 548
 549         elif self.version == "hixie-76":
 550             buf = s2b('\xff\x00')
 551             self.client.send(buf)
 552
 553         # No orderly close for 75
 554
 555     def do_handshake(self, sock, address):
 556         """
 557         do_handshake does the following:
 558         - Peek at the first few bytes from the socket.
 559         - If the connection is Flash policy request then answer it,
 560           close the socket and return.
 561         - If the connection is an HTTPS/SSL/TLS connection then SSL
 562           wrap the socket.
 563         - Read from the (possibly wrapped) socket.
 564         - If we have received a HTTP GET request and the webserver
 565           functionality is enabled, answer it, close the socket and
 566           return.
 567         - Assume we have a WebSockets connection, parse the client
 568           handshake data.
 569         - Send a WebSockets handshake server response.
 570         - Return the socket for this WebSocket client.
 571         """
 572
 573         stype = ""
 574
 575         ready = select.select([sock], [], [], 3)[0]
 576         if not ready:
 577             raise self.EClose("ignoring socket not ready")
 578         # Peek, but do not read the data so that we have a opportunity
 579         # to SSL wrap the socket first
 580         handshake = sock.recv(1024, socket.MSG_PEEK)
 581         #self.msg("Handshake [%s]" % handshake)
 582
 583         if handshake == "":
 584             raise self.EClose("ignoring empty handshake")
 585
 586         elif handshake.startswith(s2b("<policy-file-request/>")):
 587             # Answer Flash policy request
 588             handshake = sock.recv(1024)
 589             sock.send(s2b(self.policy_response))
 590             raise self.EClose("Sending flash policy response")
 591
 592         elif handshake[0] in ("\x16", "\x80", 22, 128):
 593             # SSL wrap the connection
 594             if not ssl:
 595                 raise self.EClose("SSL connection but no 'ssl' module")
 596             if not os.path.exists(self.cert):
 597                 raise self.EClose("SSL connection but '%s' not found"
 598                                   % self.cert)
 599             retsock = None
 600             try:
 601                 retsock = ssl.wrap_socket(
 602                         sock,
 603                         server_side=True,
 604                         certfile=self.cert,
 605                         keyfile=self.key)
 606             except ssl.SSLError:
 607                 _, x, _ = sys.exc_info()
 608                 if x.args[0] == ssl.SSL_ERROR_EOF:
 609                     if len(x.args) > 1:
 610                         raise self.EClose(x.args[1])
 611                     else:
 612                         raise self.EClose("Got SSL_ERROR_EOF")
 613                 else:
 614                     raise
 615
 616             scheme = "wss"
 617             stype = "SSL/TLS (wss://)"
 618
 619         elif self.ssl_only:
 620             raise self.EClose("non-SSL connection received but disallowed")
 621
 622         else:
 623             retsock = sock
 624             scheme = "ws"
 625             stype = "Plain non-SSL (ws://)"
 626
 627         wsh = WSRequestHandler(retsock, address, not self.web)
 628         if wsh.last_code == 101:
 629             # Continue on to handle WebSocket upgrade
 630             pass
 631         elif wsh.last_code == 405:
 632             raise self.EClose("Normal web request received but disallowed")
 633         elif wsh.last_code < 200 or wsh.last_code >= 300:
 634             raise self.EClose(wsh.last_message)
 635         elif self.verbose:
 636             raise self.EClose(wsh.last_message)
 637         else:
 638             raise self.EClose("")
 639
 640         h = self.headers = wsh.headers
 641         path = self.path = wsh.path
 642
 643         prot = 'WebSocket-Protocol'
 644         protocols = h.get('Sec-'+prot, h.get(prot, '')).split(',')
 645
 646         ver = h.get('Sec-WebSocket-Version')
 647         if ver:
 648             # HyBi/IETF version of the protocol
 649
 650             # HyBi-07 report version 7
 651             # HyBi-08 - HyBi-12 report version 8
 652             # HyBi-13 reports version 13
 653             if ver in ['7', '8', '13']:
 654                 self.version = "hybi-%02d" % int(ver)
 655             else:
 656                 raise self.EClose('Unsupported protocol version %s' % ver)
 657
 658             key = h['Sec-WebSocket-Key']
 659
 660             # Choose binary if client supports it
 661             if 'binary' in protocols:
 662                 self.base64 = False
 663             elif 'base64' in protocols:
 664                 self.base64 = True
 665             else:
 666                 raise self.EClose("Client must support 'binary' or 'base64' protocol")
 667
 668             # Generate the hash value for the accept header
 669             accept = b64encode(sha1(s2b(key + self.GUID)).digest())
 670
 671             response = self.server_handshake_hybi % b2s(accept)
 672             if self.base64:
 673                 response += "Sec-WebSocket-Protocol: base64\r\n"
 674             else:
 675                 response += "Sec-WebSocket-Protocol: binary\r\n"
 676             response += "\r\n"
 677
 678         else:
 679             # Hixie version of the protocol (75 or 76)
 680
 681             if h.get('key3'):
 682                 trailer = self.gen_md5(h)
 683                 pre = "Sec-"
 684                 self.version = "hixie-76"
 685             else:
 686                 trailer = ""
 687                 pre = ""
 688                 self.version = "hixie-75"
 689
 690             # We only support base64 in Hixie era
 691             self.base64 = True
 692
 693             response = self.server_handshake_hixie % (pre,
 694                     h['Origin'], pre, scheme, h['Host'], path)
 695
 696             if 'base64' in protocols:
 697                 response += "%sWebSocket-Protocol: base64\r\n" % pre
 698             else:
 699                 self.msg("Warning: client does not report 'base64' protocol support")
 700             response += "\r\n" + trailer
 701
 702         self.msg("%s: %s WebSocket connection" % (address[0], stype))
 703         self.msg("%s: Version %s, base64: '%s'" % (address[0],
 704             self.version, self.base64))
 705         if self.path != '/':
 706             self.msg("%s: Path: '%s'" % (address[0], self.path))
 707
 708
 709         # Send server WebSockets handshake response
 710         #self.msg("sending response [%s]" % response)
 711         retsock.send(s2b(response))
 712
 713         # Return the WebSockets socket which may be SSL wrapped
 714         return retsock
 715
 716
 717     #
 718     # Events that can/should be overridden in sub-classes
 719     #
 720     def started(self):
 721         """ Called after WebSockets startup """
 722         self.vmsg("WebSockets server started")
 723
 724     def poll(self):
 725         """ Run periodically while waiting for connections. """
 726         #self.vmsg("Running poll()")
 727         pass
 728
 729     def fallback_SIGCHLD(self, sig, stack):
 730         # Reap zombies when using os.fork() (python 2.4)
 731         self.vmsg("Got SIGCHLD, reaping zombies")
 732         try:
 733             result = os.waitpid(-1, os.WNOHANG)
 734             while result[0]:
 735                 self.vmsg("Reaped child process %s" % result[0])
 736                 result = os.waitpid(-1, os.WNOHANG)
 737         except (OSError):
 738             pass
 739
 740     def do_SIGINT(self, sig, stack):
 741         self.msg("Got SIGINT, exiting")
 742         sys.exit(0)
 743
 744     def top_new_client(self, startsock, address):
 745         """ Do something with a WebSockets client connection. """
 746         # Initialize per client settings
 747         self.send_parts = []
 748         self.recv_part  = None
 749         self.base64     = False
 750         self.rec        = None
 751         self.start_time = int(time.time()*1000)
 752
 753         # handler process
 754         try:
 755             try:
 756                 self.client = self.do_handshake(startsock, address)
 757
 758                 if self.record:
 759                     # Record raw frame data as JavaScript array
 760                     fname = "%s.%s" % (self.record,
 761                                         self.handler_id)
 762                     self.msg("opening record file: %s" % fname)
 763                     self.rec = open(fname, 'w+')
 764                     self.rec.write("var VNC_frame_data = [\n")
 765
 766                 self.ws_connection = True
 767                 self.new_client()
 768             except self.CClose:
 769                 # Close the client
 770                 _, exc, _ = sys.exc_info()
 771                 if self.client:
 772                     self.send_close(exc.args[0], exc.args[1])
 773             except self.EClose:
 774                 _, exc, _ = sys.exc_info()
 775                 # Connection was not a WebSockets connection
 776                 if exc.args[0]:
 777                     self.msg("%s: %s" % (address[0], exc.args[0]))
 778             except Exception:
 779                 _, exc, _ = sys.exc_info()
 780                 self.msg("handler exception: %s" % str(exc))
 781                 if self.verbose:
 782                     self.msg(traceback.format_exc())
 783         finally:
 784             if self.rec:
 785                 self.rec.write("'EOF']\n")
 786                 self.rec.close()
 787
 788             if self.client and self.client != startsock:
 789                 # Close the SSL wrapped socket
 790                 # Original socket closed by caller
 791                 self.client.close()
 792
 793     def new_client(self):
 794         """ Do something with a WebSockets client connection. """
 795         raise("WebSocketServer.new_client() must be overloaded")
 796
 797     def start_server(self):
 798         """
 799         Daemonize if requested. Listen for for connections. Run
 800         do_handshake() method for each connection. If the connection
 801         is a WebSockets client then call new_client() method (which must
 802         be overridden) for each new client connection.
 803         """
 804         lsock = self.socket(self.listen_host, self.listen_port)
 805
 806         if self.daemon:
 807             self.daemonize(keepfd=lsock.fileno(), chdir=self.web)
 808
 809         self.started()  # Some things need to happen after daemonizing
 810
 811         # Allow override of SIGINT
 812         signal.signal(signal.SIGINT, self.do_SIGINT)
 813         if not multiprocessing:
 814             # os.fork() (python 2.4) child reaper
 815             signal.signal(signal.SIGCHLD, self.fallback_SIGCHLD)
 816
 817         while True:
 818             try:
 819                 try:
 820                     self.client = None
 821                     startsock = None
 822                     pid = err = 0
 823
 824                     time_elapsed = time.time() - self.launch_time
 825                     if self.timeout and time_elapsed > self.timeout:
 826                         self.msg('listener exit due to --timeout %s'
 827                                 % self.timeout)
 828                         break
 829
 830                     try:
 831                         self.poll()
 832
 833                         ready = select.select([lsock], [], [], 1)[0]
 834                         if lsock in ready:
 835                             startsock, address = lsock.accept()
 836                         else:
 837                             continue
 838                     except Exception:
 839                         _, exc, _ = sys.exc_info()
 840                         if hasattr(exc, 'errno'):
 841                             err = exc.errno
 842                         elif hasattr(exc, 'args'):
 843                             err = exc.args[0]
 844                         else:
 845                             err = exc[0]
 846                         if err == errno.EINTR:
 847                             self.vmsg("Ignoring interrupted syscall")
 848                             continue
 849                         else:
 850                             raise
 851
 852                     if self.run_once:
 853                         # Run in same process if run_once
 854                         self.top_new_client(startsock, address)
 855                         if self.ws_connection :
 856                             self.msg('%s: exiting due to --run-once'
 857                                     % address[0])
 858                             break
 859                     elif multiprocessing:
 860                         self.vmsg('%s: new handler Process' % address[0])
 861                         p = multiprocessing.Process(
 862                                 target=self.top_new_client,
 863                                 args=(startsock, address))
 864                         p.start()
 865                         # child will not return
 866                     else:
 867                         # python 2.4
 868                         self.vmsg('%s: forking handler' % address[0])
 869                         pid = os.fork()
 870                         if pid == 0:
 871                             # child handler process
 872                             self.top_new_client(startsock, address)
 873                             break  # child process exits
 874
 875                     # parent process
 876                     self.handler_id += 1
 877
 878                 except KeyboardInterrupt:
 879                     _, exc, _ = sys.exc_info()
 880                     print("In KeyboardInterrupt")
 881                     pass
 882                 except SystemExit:
 883                     _, exc, _ = sys.exc_info()
 884                     print("In SystemExit")
 885                     break
 886                 except Exception:
 887                     _, exc, _ = sys.exc_info()
 888                     self.msg("handler exception: %s" % str(exc))
 889                     if self.verbose:
 890                         self.msg(traceback.format_exc())
 891
 892             finally:
 893                 if startsock:
 894                     startsock.close()
 895
 896
 897 # HTTP handler with WebSocket upgrade support
 898 class WSRequestHandler(SimpleHTTPRequestHandler):
 899     def __init__(self, req, addr, only_upgrade=False):
 900         self.only_upgrade = only_upgrade # only allow upgrades
 901         SimpleHTTPRequestHandler.__init__(self, req, addr, object())
 902
 903     def do_GET(self):
 904         if (self.headers.get('upgrade') and
 905                 self.headers.get('upgrade').lower() == 'websocket'):
 906
 907             if (self.headers.get('sec-websocket-key1') or
 908                     self.headers.get('websocket-key1')):
 909                 # For Hixie-76 read out the key hash
 910                 self.headers.__setitem__('key3', self.rfile.read(8))
 911
 912             # Just indicate that an WebSocket upgrade is needed
 913             self.last_code = 101
 914             self.last_message = "101 Switching Protocols"
 915         elif self.only_upgrade:
 916             # Normal web request responses are disabled
 917             self.last_code = 405
 918             self.last_message = "405 Method Not Allowed"
 919         else:
 920             SimpleHTTPRequestHandler.do_GET(self)
 921
 922     def send_response(self, code, message=None):
 923         # Save the status code
 924         self.last_code = code
 925         SimpleHTTPRequestHandler.send_response(self, code, message)
 926
 927     def log_message(self, f, *args):
 928         # Save instead of printing
 929         self.last_message = f % args
 930