python-django-websocket 0.3.0-3 / usr / share / pyshared / django_websocket / websocket.py

import collections
import select
import string
import struct
try:
    from hashlib import md5
except ImportError: #pragma NO COVER
    from md5 import md5
from errno import EINTR
from socket import error as SocketError


class MalformedWebSocket(ValueError):
    pass


def _extract_number(value):
    """
    Utility function which, given a string like 'g98sd  5[]221@1', will
    return 9852211. Used to parse the Sec-WebSocket-Key headers.
    """
    out = ""
    spaces = 0
    for char in value:
        if char in string.digits:
            out += char
        elif char == " ":
            spaces += 1
    return int(out) / spaces


def setup_websocket(request):
    if request.META.get('HTTP_CONNECTION', None) == 'Upgrade' and \
        request.META.get('HTTP_UPGRADE', None) == 'WebSocket':

        # See if they sent the new-format headers
        if 'HTTP_SEC_WEBSOCKET_KEY1' in request.META:
            protocol_version = 76
            if 'HTTP_SEC_WEBSOCKET_KEY2' not in request.META:
                raise MalformedWebSocket()
        else:
            protocol_version = 75

        # If it's new-version, we need to work out our challenge response
        if protocol_version == 76:
            key1 = _extract_number(request.META['HTTP_SEC_WEBSOCKET_KEY1'])
            key2 = _extract_number(request.META['HTTP_SEC_WEBSOCKET_KEY2'])
            # There's no content-length header in the request, but it has 8
            # bytes of data.
            key3 = request.META['wsgi.input'].read(8)
            key = struct.pack(">II", key1, key2) + key3
            handshake_response = md5(key).digest()

        location = 'ws://%s%s' % (request.get_host(), request.path)
        qs = request.META.get('QUERY_STRING')
        if qs:
            location += '?' + qs
        if protocol_version == 75:
            handshake_reply = (
                "HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
                "Upgrade: WebSocket\r\n"
                "Connection: Upgrade\r\n"
                "WebSocket-Origin: %s\r\n"
                "WebSocket-Location: %s\r\n\r\n" % (
                    request.META.get('HTTP_ORIGIN'),
                    location))
        elif protocol_version == 76:
            handshake_reply = (
                "HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
                "Upgrade: WebSocket\r\n"
                "Connection: Upgrade\r\n"
                "Sec-WebSocket-Origin: %s\r\n"
                "Sec-WebSocket-Protocol: %s\r\n"
                "Sec-WebSocket-Location: %s\r\n" % (
                    request.META.get('HTTP_ORIGIN'),
                    request.META.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'default'),
                    location))
            handshake_reply = str(handshake_reply)
            handshake_reply = '%s\r\n%s' % (handshake_reply, handshake_response)

        else:
            raise MalformedWebSocket("Unknown WebSocket protocol version.")
        socket = request.META['wsgi.input']._sock.dup()
        return WebSocket(
            socket,
            protocol=request.META.get('HTTP_WEBSOCKET_PROTOCOL'),
            version=protocol_version,
            handshake_reply=handshake_reply,
        )
    return None


class WebSocket(object):
    """
    A websocket object that handles the details of
    serialization/deserialization to the socket.

    The primary way to interact with a :class:`WebSocket` object is to
    call :meth:`send` and :meth:`wait` in order to pass messages back
    and forth with the browser.
    """
    _socket_recv_bytes = 4096


    def __init__(self, socket, protocol, version=76,
        handshake_reply=None, handshake_sent=None):
        '''
        Arguments:

        - ``socket``: An open socket that should be used for WebSocket
          communciation.
        - ``protocol``: not used yet.
        - ``version``: The WebSocket spec version to follow (default is 76)
        - ``handshake_reply``: Handshake message that should be sent to the
          client when ``send_handshake()`` is called.
        - ``handshake_sent``: Whether the handshake is already sent or not.
          Set to ``False`` to prevent ``send_handshake()`` to do anything.
        '''
        self.socket = socket
        self.protocol = protocol
        self.version = version
        self.closed = False
        self.handshake_reply = handshake_reply
        if handshake_sent is None:
            self._handshake_sent = not bool(handshake_reply)
        else:
            self._handshake_sent = handshake_sent
        self._buffer = ""
        self._message_queue = collections.deque()

    def send_handshake(self):
        self.socket.sendall(self.handshake_reply)
        self._handshake_sent = True

    @classmethod
    def _pack_message(cls, message):
        """Pack the message inside ``00`` and ``FF``

        As per the dataframing section (5.3) for the websocket spec
        """
        if isinstance(message, unicode):
            message = message.encode('utf-8')
        elif not isinstance(message, str):
            message = str(message)
        packed = "\x00%s\xFF" % message
        return packed

    def _parse_message_queue(self):
        """ Parses for messages in the buffer *buf*.  It is assumed that
        the buffer contains the start character for a message, but that it
        may contain only part of the rest of the message.

        Returns an array of messages, and the buffer remainder that
        didn't contain any full messages."""
        msgs = []
        end_idx = 0
        buf = self._buffer
        while buf:
            frame_type = ord(buf[0])
            if frame_type == 0:
                # Normal message.
                end_idx = buf.find("\xFF")
                if end_idx == -1: #pragma NO COVER
                    break
                msgs.append(buf[1:end_idx].decode('utf-8', 'replace'))
                buf = buf[end_idx+1:]
            elif frame_type == 255:
                # Closing handshake.
                assert ord(buf[1]) == 0, "Unexpected closing handshake: %r" % buf
                self.closed = True
                break
            else:
                raise ValueError("Don't understand how to parse this type of message: %r" % buf)
        self._buffer = buf
        return msgs

    def send(self, message):
        '''
        Send a message to the client. *message* should be convertable to a
        string; unicode objects should be encodable as utf-8.
        '''
        packed = self._pack_message(message)
        self.socket.sendall(packed)

    def _socket_recv(self):
        '''
        Gets new data from the socket and try to parse new messages.
        '''
        delta = self.socket.recv(self._socket_recv_bytes)
        if delta == '':
            return False
        self._buffer += delta
        msgs = self._parse_message_queue()
        self._message_queue.extend(msgs)
        return True

    def _socket_can_recv(self, timeout=0.0):
        '''
        Return ``True`` if new data can be read from the socket.
        '''
        r, w, e = [self.socket], [], []
        try:
            r, w, e = select.select(r, w, e, timeout)
        except select.error, err:
            if err.args[0] == EINTR:
                return False
            raise
        return self.socket in r

    def _get_new_messages(self):
        # read as long from socket as we need to get a new message.
        while self._socket_can_recv():
            self._socket_recv()
            if self._message_queue:
                return

    def count_messages(self):
        '''
        Returns the number of queued messages.
        '''
        self._get_new_messages()
        return len(self._message_queue)

    def has_messages(self):
        '''
        Returns ``True`` if new messages from the socket are available, else
        ``False``.
        '''
        if self._message_queue:
            return True
        self._get_new_messages()
        if self._message_queue:
            return True
        return False

    def read(self, fallback=None):
        '''
        Return new message or ``fallback`` if no message is available.
        '''
        if self.has_messages():
            return self._message_queue.popleft()
        return fallback

    def wait(self):
        '''
        Waits for and deserializes messages. Returns a single message; the
        oldest not yet processed.
        '''
        while not self._message_queue:
            # Websocket might be closed already.
            if self.closed:
                return None
            # no parsed messages, must mean buf needs more data
            new_data = self._socket_recv()
            if not new_data:
                return None
        return self._message_queue.popleft()

    def __iter__(self):
        '''
        Use ``WebSocket`` as iterator. Iteration only stops when the websocket
        gets closed by the client.
        '''
        while True:
            message = self.wait()
            if message is None:
                return
            yield message

    def _send_closing_frame(self, ignore_send_errors=False):
        '''
        Sends the closing frame to the client, if required.
        '''
        if self.version == 76 and not self.closed:
            try:
                self.socket.sendall("\xff\x00")
            except SocketError:
                # Sometimes, like when the remote side cuts off the connection,
                # we don't care about this.
                if not ignore_send_errors:
                    raise
            self.closed = True

    def close(self):
        '''
        Forcibly close the websocket.
        '''
        self._send_closing_frame()