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