Initial commit

.gitignore

@@ -0,0 +1,3 @@
+*~
+*.pyc
+build

README.md

@@ -0,0 +1,63 @@
+# python-websocket
+python-websocket is a Python/asyncore WebSocket client library.
+
+## Why would I want to use WebSocket over plain TCP for my desktop app?
+
+While the primary intention of the WebSocket protocol is to enable
+"TCP-like" full-duplex communication from within a web apps, it also
+provides a number of benefits (over plain TCP) for desktop apps:
+
+ * the WebSocket client library is responsible for message parsing,
+   and exposes a simple, event-driven API. 
+ * WebSocket connections are firewall- and proxy-friendly as they 
+   are simply upgraded HTTP connections.
+ * HTTP cookies can be attached to the WebSocket handshake, where 
+   they can be used for e.g. authentication.
+
+## Usage
+
+    def my_msg_handler(msg):
+      print 'Got "%s"!' % msg
+
+    socket = WebSocket('ws://example.com/demo', onmessage=my_msg_handler)
+    socket.onopen = lambda: socket.send('Hello world!')
+
+    try:
+      asyncore.loop()
+    except KeyboardInterrupt:
+      socket.close()
+
+## API reference
+
+### WebSocket(url, protocol=None, cookie_jar=None, onopen=None, onmessage=None, onerror=None, onclose=None)
+Returns a WebSocket connected to a remote host at the given `url`. 
+In order to allow communication over this socket, `asyncore.loop()`
+must be called by the client.
+
+The optional `protocol` parameter can be used to specify the
+sub-protocol to be used. By providing a `cookie_jar` (cookielib), 
+appropriate cookies will be sent to the server.
+
+The remaining parameters are callback functions that will be 
+invoked in the following manner:
+
+    onopen:    invoked when a connection to the remote host has been
+               successfully established
+
+    onmessage: invoked when a message is received (passing the
+               received data as an argument to the callback 
+               function)
+
+    onerror:   invoked when a communication error occured (passing 
+               an Exception instance as an argument to the callback 
+               function). If onerror is not provided, the default 
+               asyncore behavior is used (raising the exception)
+
+    onclose:   invoked when the connection has closed _normally_ (by
+               request from either client or server)
+
+### WebSocket.send(data)
+Sends `data` through the socket.
+
+### WebSocket.close()
+Closes the socket.

examples/echo.py

@@ -0,0 +1,16 @@
+#!/usr/bin/python
+import sys, os, asyncore, websocket    
+
+if __name__ == '__main__':
+    if len(sys.argv) != 2:
+        sys.stderr.write('usage: %s <message>\n' % os.path.abspath(__file__))
+        sys.exit(1)
+
+    ws = websocket.WebSocket('ws://localhost:8080', 
+                             onopen=lambda: ws.send(sys.argv[1]),
+                             onmessage=lambda m: sys.stdout.write('Echo: %s\n' % m),
+                             onclose=lambda: sys.stdout.write('Connection closed.\n'))
+    try:
+        asyncore.loop()
+    except KeyboardInterrupt:
+        ws.close()

setup.py

@@ -0,0 +1,3 @@
+from distutils.core import setup
+
+setup(name = "websocket", py_modules = ["websocket"])

websocket.py

@@ -0,0 +1,228 @@
+"""
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>
+"""
+
+import sys, re, urlparse, socket, asyncore
+
+class WebSocket(object):
+    def __init__(self, url, **kwargs):
+        self.host, self.port, self.resource, self.secure = WebSocket._parse_url(url)
+        self.protocol = kwargs.pop('protocol', None)
+        self.cookie_jar = kwargs.pop('cookie_jar', None)
+        self.onopen = kwargs.pop('onopen', None)
+        self.onmessage = kwargs.pop('onmessage', None)
+        self.onerror = kwargs.pop('onerror', None)
+        self.onclose = kwargs.pop('onclose', None)
+        if kwargs: raise ValueError('Unexpected argument(s): %s' % ', '.join(kwargs.values()))
+
+        self._dispatcher = _Dispatcher(self)
+
+    def send(self, data):
+        self._dispatcher.write('\x00' + _utf8(data) + '\xff')
+
+    def close(self):
+        self._dispatcher.handle_close()
+
+    @classmethod
+    def _parse_url(cls, url):
+        p = urlparse.urlparse(url)
+
+        if p.hostname:
+            host = p.hostname
+        else:
+            raise ValueError('URL must be absolute')
+
+        if p.fragment:
+            raise ValueError('URL must not contain a fragment component')
+
+        if p.scheme == 'ws':
+            secure = False
+            port = p.port or 80
+        elif p.scheme == 'wss':
+            raise NotImplementedError('Secure WebSocket not yet supported')
+            # secure = True
+            # port = p.port or 443
+        else:
+            raise ValueError('Invalid URL scheme')
+
+        resource = p.path or u'/'
+        if p.query: resource += u'?' + p.query
+        return (host, port, resource, secure)
+
+    #@classmethod
+    #def _generate_key(cls):
+    #    spaces = random.randint(1, 12)
+    #    number = random.randint(0, 0xffffffff/spaces)
+    #    key = list(str(number*spaces))
+    #    chars = map(unichr, range(0x21, 0x2f) + range(0x3a, 0x7e))
+    #    random_inserts = random.sample(xrange(len(key)), random.randint(1,12))
+    #    for (i, c) in [(r, random.choice(chars)) for r in random_inserts]:
+    #        key.insert(i, c)
+    #    print key
+    #    return ''.join(key)
+
+class WebSocketError(Exception):
+    def _init_(self, value):
+        self.value = value
+
+    def _str_(self):
+        return str(self.value)
+
+class _Dispatcher(asyncore.dispatcher):
+    def __init__(self, ws):
+        asyncore.dispatcher.__init__(self)
+        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.connect((ws.host, ws.port))
+
+        self.ws = ws
+        self._read_buffer = ''
+        self._write_buffer = ''
+        self._handshake_complete = False
+
+        if self.ws.port != 80:
+            hostport = '%s:%d' % (self.ws.host, self.ws.port)
+        else:
+            hostport = self.ws.host
+
+        fields = [
+            'Upgrade: WebSocket',
+            'Connection: Upgrade',
+            'Host: ' + hostport,
+            'Origin: http://' + hostport,
+            #'Sec-WebSocket-Key1: %s' % WebSocket.generate_key(),
+            #'Sec-WebSocket-Key2: %s' % WebSocket.generate_key()
+        ]
+        if self.ws.protocol: fields['Sec-WebSocket-Protocol'] = self.ws.protocol
+        if self.ws.cookie_jar:
+            cookies = filter(lambda c: _cookie_for_domain(c, _eff_host(self.ws.host)) and \
+                             _cookie_for_path(c, self.ws.resource) and \
+                             not c.is_expired(), self.ws.cookie_jar)
+
+            for cookie in cookies:
+                fields.append('Cookie: %s=%s' % (cookie.name, cookie.value))
+
+        # key3 = ''.join(map(unichr, (random.randrange(256) for i in xrange(8))))
+        self.write(_utf8('GET %s HTTP/1.1\r\n' \
+                         '%s\r\n\r\n' % (self.ws.resource,
+                                         '\r\n'.join(fields))))
+                                         # key3)))
+
+    def handl_expt(self):
+        self.handle_error()
+
+    def handle_error(self):
+        self.close()
+        t, e, trace = sys.exc_info()
+        if self.ws.onerror:
+            self.ws.onerror(e)
+        else:
+            asyncore.dispatcher.handle_error(self)
+
+    def handle_close(self):
+        self.close()
+        if self.ws.onclose:
+            self.ws.onclose()
+
+    def handle_read(self):
+        if self._handshake_complete:
+            self._read_until('\xff', self._handle_frame)
+        else:
+            self._read_until('\r\n\r\n', self._handle_header)
+
+    def handle_write(self):
+        sent = self.send(self._write_buffer)
+        self._write_buffer = self._write_buffer[sent:]
+
+    def writable(self):
+        return len(self._write_buffer) > 0
+
+    def write(self, data):
+        self._write_buffer += data # TODO: separate buffer for handshake from data to
+                                   # prevent mix-up when send() is called before
+                                   # handshake is complete?
+
+    def _read_until(self, delimiter, callback):
+        self._read_buffer += self.recv(4096)
+        pos = self._read_buffer.find(delimiter)+len(delimiter)+1
+        if pos > 0:
+            data = self._read_buffer[:pos]
+            self._read_buffer = self._read_buffer[pos:]
+            if data:
+                callback(data)
+
+    def _handle_frame(self, frame):
+        assert frame[-1] == '\xff'
+        if frame[0] != '\x00':
+            raise WebSocketError('WebSocket stream error')
+
+        if self.ws.onmessage:
+            self.ws.onmessage(frame[1:-1])
+        # TODO: else raise WebSocketError('No message handler defined')
+
+    def _handle_header(self, header):
+        assert header[-4:] == '\r\n\r\n'
+        start_line, fields = _parse_http_header(header)
+        if start_line != 'HTTP/1.1 101 Web Socket Protocol Handshake' or \
+           fields.get('Connection', None) != 'Upgrade' or \
+           fields.get('Upgrade', None) != 'WebSocket':
+            raise WebSocketError('Invalid server handshake')
+
+        self._handshake_complete = True
+        if self.ws.onopen:
+            self.ws.onopen()
+
+_IPV4_RE = re.compile(r'\.\d+$')
+_PATH_SEP = re.compile(r'/+')
+
+def _parse_http_header(header):
+    def split_field(field):
+        k, v = field.split(':', 1)
+        return (k, v.strip())
+
+    lines = header.strip().split('\r\n')
+    if len(lines) > 0:
+        start_line = lines[0]
+    else:
+        start_line = None
+
+    return (start_line, dict(map(split_field, lines[1:])))
+
+def _eff_host(host):
+    if host.find('.') == -1 and not _IPV4_RE.search(host):
+        return host + '.local'
+    return host
+
+def _cookie_for_path(cookie, path):
+    if not cookie.path or path == '' or path == '/':
+        return True
+    path = _PATH_SEP.split(path)[1:]
+    cookie_path = _PATH_SEP.split(cookie.path)[1:]
+    for p1, p2 in map(lambda *ps: ps, path, cookie_path):
+        if p1 == None:
+            return True
+        elif p1 != p2:
+            return False
+
+    return True
+
+def _cookie_for_domain(cookie, domain):
+    if not cookie.domain:
+        return True
+    elif cookie.domain[0] == '.':
+        return domain.endswith(cookie.domain)
+    else:
+        return cookie.domain == domain
+
+def _utf8(s):
+    return s.encode('utf-8')