Statistics
| Branch: | Tag: | Revision:

one / src / sunstone / public / vendor / noVNC / utils / websocket.py @ 5fcb9de8

History | View | Annotate | Download (14.4 KB)

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