utils/websocket.py

   1 #!/usr/bin/env 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             try:
 151                 var, val = line.split(": ")
 152             except:
 153                 raise Exception("Invalid handshake header: %s" % line)
 154             ret[var] = val
 155
 156         if req_lines[-2] == "":
 157             ret['key3'] = req_lines[-1]
 158
 159         return ret
 160
 161     @staticmethod
 162     def gen_md5(keys):
 163         """ Generate hash value for WebSockets handshake v76. """
 164         key1 = keys['Sec-WebSocket-Key1']
 165         key2 = keys['Sec-WebSocket-Key2']
 166         key3 = keys['key3']
 167         spaces1 = key1.count(" ")
 168         spaces2 = key2.count(" ")
 169         num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
 170         num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
 171
 172         return md5(struct.pack('>II8s', num1, num2, key3)).digest()
 173
 174
 175     #
 176     # WebSocketServer logging/output functions
 177     #
 178
 179     def traffic(self, token="."):
 180         """ Show traffic flow in verbose mode. """
 181         if self.verbose and not self.daemon:
 182             sys.stdout.write(token)
 183             sys.stdout.flush()
 184
 185     def msg(self, msg):
 186         """ Output message with handler_id prefix. """
 187         if not self.daemon:
 188             print "% 3d: %s" % (self.handler_id, msg)
 189
 190     def vmsg(self, msg):
 191         """ Same as msg() but only if verbose. """
 192         if self.verbose:
 193             self.msg(msg)
 194
 195     #
 196     # Main WebSocketServer methods
 197     #
 198
 199     def do_handshake(self, sock, address):
 200         """
 201         do_handshake does the following:
 202         - Peek at the first few bytes from the socket.
 203         - If the connection is Flash policy request then answer it,
 204           close the socket and return.
 205         - If the connection is an HTTPS/SSL/TLS connection then SSL
 206           wrap the socket.
 207         - Read from the (possibly wrapped) socket.
 208         - If we have received a HTTP GET request and the webserver
 209           functionality is enabled, answer it, close the socket and
 210           return.
 211         - Assume we have a WebSockets connection, parse the client
 212           handshake data.
 213         - Send a WebSockets handshake server response.
 214         - Return the socket for this WebSocket client.
 215         """
 216
 217         stype = ""
 218
 219         ready = select.select([sock], [], [], 3)[0]
 220         if not ready:
 221             raise self.EClose("ignoring socket not ready")
 222         # Peek, but do not read the data so that we have a opportunity
 223         # to SSL wrap the socket first
 224         handshake = sock.recv(1024, socket.MSG_PEEK)
 225         #self.msg("Handshake [%s]" % repr(handshake))
 226
 227         if handshake == "":
 228             raise self.EClose("ignoring empty handshake")
 229
 230         elif handshake.startswith("<policy-file-request/>"):
 231             # Answer Flash policy request
 232             handshake = sock.recv(1024)
 233             sock.send(self.policy_response)
 234             raise self.EClose("Sending flash policy response")
 235
 236         elif handshake[0] in ("\x16", "\x80"):
 237             # SSL wrap the connection
 238             if not os.path.exists(self.cert):
 239                 raise self.EClose("SSL connection but '%s' not found"
 240                                   % self.cert)
 241             try:
 242                 retsock = ssl.wrap_socket(
 243                         sock,
 244                         server_side=True,
 245                         certfile=self.cert,
 246                         keyfile=self.key)
 247             except ssl.SSLError, x:
 248                 if x.args[0] == ssl.SSL_ERROR_EOF:
 249                     raise self.EClose("")
 250                 else:
 251                     raise
 252
 253             scheme = "wss"
 254             stype = "SSL/TLS (wss://)"
 255
 256         elif self.ssl_only:
 257             raise self.EClose("non-SSL connection received but disallowed")
 258
 259         else:
 260             retsock = sock
 261             scheme = "ws"
 262             stype = "Plain non-SSL (ws://)"
 263
 264         # Now get the data from the socket
 265         handshake = retsock.recv(4096)
 266
 267         if len(handshake) == 0:
 268             raise self.EClose("Client closed during handshake")
 269
 270         # Check for and handle normal web requests
 271         if handshake.startswith('GET ') and \
 272             handshake.find('Upgrade: WebSocket\r\n') == -1:
 273             if not self.web:
 274                 raise self.EClose("Normal web request received but disallowed")
 275             sh = SplitHTTPHandler(handshake, retsock, address)
 276             if sh.last_code < 200 or sh.last_code >= 300:
 277                 raise self.EClose(sh.last_message)
 278             elif self.verbose:
 279                 raise self.EClose(sh.last_message)
 280             else:
 281                 raise self.EClose("")
 282
 283         #self.msg("handshake: " + repr(handshake))
 284         # Parse client WebSockets handshake
 285         h = self.parse_handshake(handshake)
 286
 287         if h.get('key3'):
 288             trailer = self.gen_md5(h)
 289             pre = "Sec-"
 290             ver = 76
 291         else:
 292             trailer = ""
 293             pre = ""
 294             ver = 75
 295
 296         self.msg("%s: %s WebSocket connection (version %s)"
 297                     % (address[0], stype, ver))
 298
 299         # Send server WebSockets handshake response
 300         response = self.server_handshake % (pre, h['Origin'], pre,
 301                 scheme, h['Host'], h['path'], pre, trailer)
 302         #self.msg("sending response:", repr(response))
 303         retsock.send(response)
 304
 305         # Return the WebSockets socket which may be SSL wrapped
 306         return retsock
 307
 308
 309     #
 310     # Events that can/should be overridden in sub-classes
 311     #
 312     def started(self):
 313         """ Called after WebSockets startup """
 314         self.vmsg("WebSockets server started")
 315
 316     def poll(self):
 317         """ Run periodically while waiting for connections. """
 318         #self.vmsg("Running poll()")
 319         pass
 320
 321     def top_SIGCHLD(self, sig, stack):
 322         # Reap zombies after calling child SIGCHLD handler
 323         self.do_SIGCHLD(sig, stack)
 324         self.vmsg("Got SIGCHLD, reaping zombies")
 325         try:
 326             result = os.waitpid(-1, os.WNOHANG)
 327             while result[0]:
 328                 self.vmsg("Reaped child process %s" % result[0])
 329                 result = os.waitpid(-1, os.WNOHANG)
 330         except (OSError):
 331             pass
 332
 333     def do_SIGCHLD(self, sig, stack):
 334         pass
 335
 336     def do_SIGINT(self, sig, stack):
 337         self.msg("Got SIGINT, exiting")
 338         sys.exit(0)
 339
 340     def new_client(self, client):
 341         """ Do something with a WebSockets client connection. """
 342         raise("WebSocketServer.new_client() must be overloaded")
 343
 344     def start_server(self):
 345         """
 346         Daemonize if requested. Listen for for connections. Run
 347         do_handshake() method for each connection. If the connection
 348         is a WebSockets client then call new_client() method (which must
 349         be overridden) for each new client connection.
 350         """
 351
 352         lsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 353         lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 354         lsock.bind((self.listen_host, self.listen_port))
 355         lsock.listen(100)
 356
 357         if self.daemon:
 358             self.daemonize(self, keepfd=lsock.fileno())
 359
 360         self.started()  # Some things need to happen after daemonizing
 361
 362         # Reep zombies
 363         signal.signal(signal.SIGCHLD, self.top_SIGCHLD)
 364         signal.signal(signal.SIGINT, self.do_SIGINT)
 365
 366         while True:
 367             try:
 368                 try:
 369                     csock = startsock = None
 370                     pid = err = 0
 371
 372                     try:
 373                         self.poll()
 374
 375                         ready = select.select([lsock], [], [], 1)[0];
 376                         if lsock in ready:
 377                             startsock, address = lsock.accept()
 378                         else:
 379                             continue
 380                     except Exception, exc:
 381                         if hasattr(exc, 'errno'):
 382                             err = exc.errno
 383                         else:
 384                             err = exc[0]
 385                         if err == errno.EINTR:
 386                             self.vmsg("Ignoring interrupted syscall")
 387                             continue
 388                         else:
 389                             raise
 390
 391                     self.vmsg('%s: forking handler' % address[0])
 392                     pid = os.fork()
 393
 394                     if pid == 0:
 395                         # handler process
 396                         csock = self.do_handshake(startsock, address)
 397                         self.new_client(csock)
 398                     else:
 399                         # parent process
 400                         self.handler_id += 1
 401
 402                 except self.EClose, exc:
 403                     # Connection was not a WebSockets connection
 404                     if exc.args[0]:
 405                         self.msg("%s: %s" % (address[0], exc.args[0]))
 406                 except KeyboardInterrupt, exc:
 407                     pass
 408                 except Exception, exc:
 409                     self.msg("handler exception: %s" % str(exc))
 410                     if self.verbose:
 411                         self.msg(traceback.format_exc())
 412
 413             finally:
 414                 if csock and csock != startsock:
 415                     csock.close()
 416                 if startsock:
 417                     startsock.close()
 418
 419             if pid == 0:
 420                 break # Child process exits
 421
 422
 423 # HTTP handler with request from a string and response to a socket
 424 class SplitHTTPHandler(SimpleHTTPRequestHandler):
 425     def __init__(self, req, resp, addr):
 426         # Save the response socket
 427         self.response = resp
 428         SimpleHTTPRequestHandler.__init__(self, req, addr, object())
 429
 430     def setup(self):
 431         self.connection = self.response
 432         # Duck type request string to file object
 433         self.rfile = StringIO(self.request)
 434         self.wfile = self.connection.makefile('wb', self.wbufsize)
 435
 436     def send_response(self, code, message=None):
 437         # Save the status code
 438         self.last_code = code
 439         SimpleHTTPRequestHandler.send_response(self, code, message)
 440
 441     def log_message(self, f, *args):
 442         # Save instead of printing
 443         self.last_message = f % args
 444
 445