utils/websocket.py

   1 #!/usr/bin/python
   2
   3 '''
   4 Python WebSocket library with support for "wss://" encryption.
   5 Copyright 2010 Joel Martin
   6 Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
   7
   8 You can make a cert/key with openssl using:
   9 openssl req -new -x509 -days 365 -nodes -out self.pem -keyout self.pem
  10 as taken from http://docs.python.org/dev/library/ssl.html#certificates
  11
  12 '''
  13
  14 import sys, socket, ssl, struct, traceback, select
  15 import os, resource, errno, signal # daemonizing
  16 from SimpleHTTPServer import SimpleHTTPRequestHandler
  17 from cStringIO import StringIO
  18 from base64 import b64encode, b64decode
  19 try:
  20     from hashlib import md5
  21 except:
  22     from md5 import md5  # Support python 2.4
  23 from urlparse import urlsplit
  24 from cgi import parse_qsl
  25
  26 class WebSocketServer(object):
  27     """
  28     WebSockets server class.
  29     Must be sub-classed with new_client method definition.
  30     """
  31
  32     server_handshake = """HTTP/1.1 101 Web Socket Protocol Handshake\r
  33 Upgrade: WebSocket\r
  34 Connection: Upgrade\r
  35 %sWebSocket-Origin: %s\r
  36 %sWebSocket-Location: %s://%s%s\r
  37 %sWebSocket-Protocol: sample\r
  38 \r
  39 %s"""
  40
  41     policy_response = """<cross-domain-policy><allow-access-from domain="*" to-ports="*" /></cross-domain-policy>\n"""
  42
  43     class EClose(Exception):
  44         pass
  45
  46     def __init__(self, listen_host='', listen_port=None,
  47             verbose=False, cert='', key='', ssl_only=None,
  48             daemon=False, record='', web=''):
  49
  50         # settings
  51         self.verbose     = verbose
  52         self.listen_host = listen_host
  53         self.listen_port = listen_port
  54         self.ssl_only    = ssl_only
  55         self.daemon      = daemon
  56
  57
  58         # Make paths settings absolute
  59         self.cert = os.path.abspath(cert)
  60         self.key = self.web = self.record = ''
  61         if key:
  62             self.key = os.path.abspath(key)
  63         if web:
  64             self.web = os.path.abspath(web)
  65         if record:
  66             self.record = os.path.abspath(record)
  67
  68         if self.web:
  69             os.chdir(self.web)
  70
  71         self.handler_id  = 1
  72
  73         print "WebSocket server settings:"
  74         print "  - Listen on %s:%s" % (
  75                 self.listen_host, self.listen_port)
  76         print "  - Flash security policy server"
  77         if self.web:
  78             print "  - Web server"
  79         if os.path.exists(self.cert):
  80             print "  - SSL/TLS support"
  81             if self.ssl_only:
  82                 print "  - Deny non-SSL/TLS connections"
  83         else:
  84             print "  - No SSL/TLS support (no cert file)"
  85         if self.daemon:
  86             print "  - Backgrounding (daemon)"
  87
  88     #
  89     # WebSocketServer static methods
  90     #
  91     @staticmethod
  92     def daemonize(self, keepfd=None):
  93         os.umask(0)
  94         if self.web:
  95             os.chdir(self.web)
  96         else:
  97             os.chdir('/')
  98         os.setgid(os.getgid())  # relinquish elevations
  99         os.setuid(os.getuid())  # relinquish elevations
 100
 101         # Double fork to daemonize
 102         if os.fork() > 0: os._exit(0)  # Parent exits
 103         os.setsid()                    # Obtain new process group
 104         if os.fork() > 0: os._exit(0)  # Parent exits
 105
 106         # Signal handling
 107         def terminate(a,b): os._exit(0)
 108         signal.signal(signal.SIGTERM, terminate)
 109         signal.signal(signal.SIGINT, signal.SIG_IGN)
 110
 111         # Close open files
 112         maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
 113         if maxfd == resource.RLIM_INFINITY: maxfd = 256
 114         for fd in reversed(range(maxfd)):
 115             try:
 116                 if fd != keepfd:
 117                     os.close(fd)
 118             except OSError, exc:
 119                 if exc.errno != errno.EBADF: raise
 120
 121         # Redirect I/O to /dev/null
 122         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdin.fileno())
 123         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdout.fileno())
 124         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno())
 125
 126     @staticmethod
 127     def encode(buf):
 128         """ Encode a WebSocket packet. """
 129         buf = b64encode(buf)
 130         return "\x00%s\xff" % buf
 131
 132     @staticmethod
 133     def decode(buf):
 134         """ Decode WebSocket packets. """
 135         if buf.count('\xff') > 1:
 136             return [b64decode(d[1:]) for d in buf.split('\xff')]
 137         else:
 138             return [b64decode(buf[1:-1])]
 139
 140     @staticmethod
 141     def parse_handshake(handshake):
 142         """ Parse fields from client WebSockets handshake. """
 143         ret = {}
 144         req_lines = handshake.split("\r\n")
 145         if not req_lines[0].startswith("GET "):
 146             raise Exception("Invalid handshake: no GET request line")
 147         ret['path'] = req_lines[0].split(" ")[1]
 148         for line in req_lines[1:]:
 149             if line == "": break
 150             var, val = line.split(": ")
 151             ret[var] = val
 152
 153         if req_lines[-2] == "":
 154             ret['key3'] = req_lines[-1]
 155
 156         return ret
 157
 158     @staticmethod
 159     def gen_md5(keys):
 160         """ Generate hash value for WebSockets handshake v76. """
 161         key1 = keys['Sec-WebSocket-Key1']
 162         key2 = keys['Sec-WebSocket-Key2']
 163         key3 = keys['key3']
 164         spaces1 = key1.count(" ")
 165         spaces2 = key2.count(" ")
 166         num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
 167         num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
 168
 169         return md5(struct.pack('>II8s', num1, num2, key3)).digest()
 170
 171
 172     #
 173     # WebSocketServer logging/output functions
 174     #
 175
 176     def traffic(self, token="."):
 177         """ Show traffic flow in verbose mode. """
 178         if self.verbose and not self.daemon:
 179             sys.stdout.write(token)
 180             sys.stdout.flush()
 181
 182     def msg(self, msg):
 183         """ Output message with handler_id prefix. """
 184         if not self.daemon:
 185             print "% 3d: %s" % (self.handler_id, msg)
 186
 187     def vmsg(self, msg):
 188         """ Same as msg() but only if verbose. """
 189         if self.verbose:
 190             self.msg(msg)
 191
 192     #
 193     # Main WebSocketServer methods
 194     #
 195
 196     def do_handshake(self, sock, address):
 197         """
 198         do_handshake does the following:
 199         - Peek at the first few bytes from the socket.
 200         - If the connection is Flash policy request then answer it,
 201           close the socket and return.
 202         - If the connection is an HTTPS/SSL/TLS connection then SSL
 203           wrap the socket.
 204         - Read from the (possibly wrapped) socket.
 205         - If we have received a HTTP GET request and the webserver
 206           functionality is enabled, answer it, close the socket and
 207           return.
 208         - Assume we have a WebSockets connection, parse the client
 209           handshake data.
 210         - Send a WebSockets handshake server response.
 211         - Return the socket for this WebSocket client.
 212         """
 213
 214         stype = ""
 215
 216         # Peek, but don't read the data
 217         handshake = sock.recv(1024, socket.MSG_PEEK)
 218         #self.msg("Handshake [%s]" % repr(handshake))
 219
 220         if handshake == "":
 221             raise self.EClose("ignoring empty handshake")
 222
 223         elif handshake.startswith("<policy-file-request/>"):
 224             # Answer Flash policy request
 225             handshake = sock.recv(1024)
 226             sock.send(self.policy_response)
 227             raise self.EClose("Sending flash policy response")
 228
 229         elif handshake[0] in ("\x16", "\x80"):
 230             # SSL wrap the connection
 231             if not os.path.exists(self.cert):
 232                 raise self.EClose("SSL connection but '%s' not found"
 233                                   % self.cert)
 234             try:
 235                 retsock = ssl.wrap_socket(
 236                         sock,
 237                         server_side=True,
 238                         certfile=self.cert,
 239                         keyfile=self.key)
 240             except ssl.SSLError, x:
 241                 if x.args[0] == ssl.SSL_ERROR_EOF:
 242                     raise self.EClose("")
 243                 else:
 244                     raise
 245
 246             scheme = "wss"
 247             stype = "SSL/TLS (wss://)"
 248
 249         elif self.ssl_only:
 250             raise self.EClose("non-SSL connection received but disallowed")
 251
 252         else:
 253             retsock = sock
 254             scheme = "ws"
 255             stype = "Plain non-SSL (ws://)"
 256
 257         # Now get the data from the socket
 258         handshake = retsock.recv(4096)
 259         #self.msg("handshake: " + repr(handshake))
 260
 261         if len(handshake) == 0:
 262             raise self.EClose("Client closed during handshake")
 263
 264         # Check for and handle normal web requests
 265         if handshake.startswith('GET ') and \
 266             handshake.find('Upgrade: WebSocket\r\n') == -1:
 267             if not self.web:
 268                 raise self.EClose("Normal web request received but disallowed")
 269             sh = SplitHTTPHandler(handshake, retsock, address)
 270             if sh.last_code < 200 or sh.last_code >= 300:
 271                 raise self.EClose(sh.last_message)
 272             elif self.verbose:
 273                 raise self.EClose(sh.last_message)
 274             else:
 275                 raise self.EClose("")
 276
 277         # Parse client WebSockets handshake
 278         h = self.parse_handshake(handshake)
 279
 280         if h.get('key3'):
 281             trailer = self.gen_md5(h)
 282             pre = "Sec-"
 283             ver = 76
 284         else:
 285             trailer = ""
 286             pre = ""
 287             ver = 75
 288
 289         self.msg("%s: %s WebSocket connection (version %s)"
 290                     % (address[0], stype, ver))
 291
 292         # Send server WebSockets handshake response
 293         response = self.server_handshake % (pre, h['Origin'], pre,
 294                 scheme, h['Host'], h['path'], pre, trailer)
 295         #self.msg("sending response:", repr(response))
 296         retsock.send(response)
 297
 298         # Return the WebSockets socket which may be SSL wrapped
 299         return retsock
 300
 301
 302     #
 303     # Events that can/should be overridden in sub-classes
 304     #
 305     def started(self):
 306         """ Called after WebSockets startup """
 307         self.vmsg("WebSockets server started")
 308
 309     def poll(self):
 310         """ Run periodically while waiting for connections. """
 311         self.msg("Running poll()")
 312
 313     def do_SIGCHLD(self, sig, stack):
 314         self.vmsg("Got SIGCHLD, ignoring")
 315
 316     def do_SIGINT(self, sig, stack):
 317         self.msg("Got SIGINT, exiting")
 318         sys.exit(0)
 319
 320     def new_client(self, client):
 321         """ Do something with a WebSockets client connection. """
 322         raise("WebSocketServer.new_client() must be overloaded")
 323
 324     def start_server(self):
 325         """
 326         Daemonize if requested. Listen for for connections. Run
 327         do_handshake() method for each connection. If the connection
 328         is a WebSockets client then call new_client() method (which must
 329         be overridden) for each new client connection.
 330         """
 331
 332         lsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 333         lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 334         lsock.bind((self.listen_host, self.listen_port))
 335         lsock.listen(100)
 336
 337         if self.daemon:
 338             self.daemonize(self, keepfd=lsock.fileno())
 339
 340         self.started()  # Some things need to happen after daemonizing
 341
 342         # Reep zombies
 343         signal.signal(signal.SIGCHLD, self.do_SIGCHLD)
 344         signal.signal(signal.SIGINT, self.do_SIGINT)
 345
 346         while True:
 347             try:
 348                 try:
 349                     csock = startsock = None
 350                     pid = err = 0
 351
 352                     try:
 353                         self.poll()
 354
 355                         ready = select.select([lsock], [], [], 1)[0];
 356                         if lsock in ready:
 357                             startsock, address = lsock.accept()
 358                         else:
 359                             continue
 360                     except Exception, exc:
 361                         if hasattr(exc, 'errno'):
 362                             err = exc.errno
 363                         else:
 364                             err = exc[0]
 365                         if err == errno.EINTR:
 366                             self.vmsg("Ignoring interrupted syscall")
 367                             continue
 368                         else:
 369                             raise
 370
 371                     self.vmsg('%s: forking handler' % address[0])
 372                     pid = os.fork()
 373
 374                     if pid == 0:
 375                         # handler process
 376                         csock = self.do_handshake(startsock, address)
 377                         self.new_client(csock)
 378                     else:
 379                         # parent process
 380                         self.handler_id += 1
 381
 382                 except self.EClose, exc:
 383                     # Connection was not a WebSockets connection
 384                     if exc.args[0]:
 385                         self.msg("%s: %s" % (address[0], exc.args[0]))
 386                 except KeyboardInterrupt, exc:
 387                     pass
 388                 except Exception, exc:
 389                     self.msg("handler exception: %s" % str(exc))
 390                     if self.verbose:
 391                         self.msg(traceback.format_exc())
 392
 393             finally:
 394                 if csock and csock != startsock:
 395                     csock.close()
 396                 if startsock:
 397                     startsock.close()
 398
 399             if pid == 0:
 400                 break # Child process exits
 401
 402
 403 # HTTP handler with request from a string and response to a socket
 404 class SplitHTTPHandler(SimpleHTTPRequestHandler):
 405     def __init__(self, req, resp, addr):
 406         # Save the response socket
 407         self.response = resp
 408         SimpleHTTPRequestHandler.__init__(self, req, addr, object())
 409
 410     def setup(self):
 411         self.connection = self.response
 412         # Duck type request string to file object
 413         self.rfile = StringIO(self.request)
 414         self.wfile = self.connection.makefile('wb', self.wbufsize)
 415
 416     def send_response(self, code, message=None):
 417         # Save the status code
 418         self.last_code = code
 419         SimpleHTTPRequestHandler.send_response(self, code, message)
 420
 421     def log_message(self, f, *args):
 422         # Save instead of printing
 423         self.last_message = f % args
 424
 425