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