websocket.py

# -*- test-case-name: twisted.web.test.test_websocket -*-
# Copyright (c) 2009 Twisted Matrix Laboratories.
# See LICENSE for details.

"""
Note: This is from the associated branch for http://twistedmatrix.com/trac/ticket/4173
and includes support for the hixie-76 handshake.

WebSocket server protocol.

See U{http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol} for the
current version of the specification.

@since: 10.1
"""

import base64
from hashlib import md5, sha1
import itertools
import re
import struct

from twisted.internet import interfaces
from twisted.python import log
from twisted.web._newclient import makeStatefulDispatcher
from twisted.protocols.policies import WrappingFactory, ProtocolWrapper
from twisted.web.http import datetimeToString
from twisted.web.http import _IdentityTransferDecoder
from twisted.web.server import Request, Site, version, unquote
from zope.interface import implements


_ascii_numbers = frozenset(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])


(OPCODE_CONT, OPCODE_TEXT, OPCODE_BINARY,
 OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG) = (0x0, 0x1, 0x2, 0x8, 0x9, 0xA)


ALL_OPCODES = (OPCODE_CONT, OPCODE_TEXT, OPCODE_BINARY,
               OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG)


CONTROL_OPCODES = (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG)
DATA_OPCODES = (OPCODE_TEXT, OPCODE_BINARY)


class WebSocketRequest(Request):
    """
    A general purpose L{Request} supporting connection upgrade for WebSocket.
    """

    ACCEPT_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

    def process(self):
        # get upgrade headers and switch them to lower case
        upgrade_headers = self.requestHeaders.getRawHeaders("Upgrade") or []
        upgrade_headers = [h.lower() for h in upgrade_headers]

        connection_headers = self.requestHeaders.getRawHeaders("Connection") or []
        # get all connection_headers, split each at ',',
        # join into a single list and switch them to lower case
        connection_headers = itertools.chain(*[re.split(r',\s*', h) for h in connection_headers])
        connection_headers = [h.lower() for h in connection_headers]

        if ("websocket" in upgrade_headers and "upgrade" in connection_headers):
            return self.processWebSocket()
        else:
            return Request.process(self)

    def processWebSocket(self):
        """
        Process a specific web socket request.
        """
        # get site from channel
        self.site = self.channel.site

        # set various default headers
        self.setHeader("server", version)
        self.setHeader("date", datetimeToString())

        # Resource Identification
        self.prepath = []
        self.postpath = map(unquote, self.path[1:].split("/"))
        self.renderWebSocket()


    def _clientHandshake76(self):
        """
        Complete hixie-76 handshake, which consists of a challenge and response.

        If the request is not identified with a proper WebSocket handshake, the
        connection will be closed. Otherwise, the response to the handshake is
        sent and a C{WebSocketHandler} is created to handle the request.
        """
        def finish():
            self.channel.transport.loseConnection()
        if self.queued:
            return finish()

        secKey1 = self.requestHeaders.getRawHeaders("Sec-WebSocket-Key1", [])
        secKey2 = self.requestHeaders.getRawHeaders("Sec-WebSocket-Key2", [])

        if len(secKey1) != 1 or len(secKey2) != 1:
            return finish()

        # copied
        originHeaders = self.requestHeaders.getRawHeaders("Origin", [])
        if len(originHeaders) != 1:
            return finish()
        hostHeaders = self.requestHeaders.getRawHeaders("Host", [])
        if len(hostHeaders) != 1:
            return finish()
        handlerFactory = self.site.handlers.get(self.uri)
        if not handlerFactory:
            return finish()

        # key1 and key2 exist and are a string of characters
        # filter both keys to get a string with all numbers in order
        key1 = secKey1[0]
        key2 = secKey2[0]
        numBuffer1 = ''.join([x for x in key1 if x in _ascii_numbers])
        numBuffer2 = ''.join([x for x in key2 if x in _ascii_numbers])

        # make sure numbers actually exist
        if not numBuffer1 or not numBuffer2:
            return finish()

        # these should be int-like
        num1 = int(numBuffer1)
        num2 = int(numBuffer2)

        # count the number of spaces in each character string
        numSpaces1 = 0
        for x in key1:
            if x == ' ':
                numSpaces1 += 1
        numSpaces2 = 0
        for x in key2:
            if x == ' ':
                numSpaces2 += 1

        # there should be at least one space in each
        if numSpaces1 == 0 or numSpaces2 == 0:
            return finish()

        # get two resulting numbers, as specified in hixie-76
        num1 = num1 / numSpaces1
        num2 = num2 / numSpaces2

        transport = WebSocketTransport(self)
        handler = handlerFactory(transport)
        transport._attachHandler(handler)

        self.channel.setRawMode()

        def finishHandshake(nonce):
            """ Receive nonce value from request body, and calculate repsonse. """
            protocolHeaders = self.requestHeaders.getRawHeaders(
                "WebSocket-Protocol", [])
            if len(protocolHeaders) not in (0,  1):
                return finish()
            if protocolHeaders:
                if protocolHeaders[0] not in self.site.supportedProtocols:
                    return finish()
                protocolHeader = protocolHeaders[0]
            else:
                protocolHeader = None

            originHeader = originHeaders[0]
            hostHeader = hostHeaders[0]
            self.startedWriting = True
            handshake = [
                "HTTP/1.1 101 Web Socket Protocol Handshake",
                "Upgrade: WebSocket",
                "Connection: Upgrade"]
            handshake.append("Sec-WebSocket-Origin: %s" % (originHeader))
            if self.isSecure():
                scheme = "wss"
            else:
                scheme = "ws"
            handshake.append(
                "Sec-WebSocket-Location: %s://%s%s" % (
                scheme, hostHeader, self.uri))

            if protocolHeader is not None:
                handshake.append("Sec-WebSocket-Protocol: %s" % protocolHeader)

            for header in handshake:
                self.write("%s\r\n" % header)

            self.write("\r\n")

            # concatenate num1 (32 bit in), num2 (32 bit int), nonce, and take md5 of result
            res = struct.pack('>II8s', num1, num2, nonce)
            server_response = md5(res).digest()
            self.write(server_response)

            # XXX we probably don't want to set _transferDecoder
            self.channel._transferDecoder = WebSocketFrameDecoder(
                self, handler)

            transport._connectionMade()

        # we need the nonce from the request body
        self.channel._transferDecoder = _IdentityTransferDecoder(0, lambda _ : None, finishHandshake)


    def _checkClientHandshake(self):
        """
        Verify client handshake, closing the connection in case of problem.

        @return: C{None} if a problem was detected, or a tuple of I{Origin}
            header, I{Host} header, I{WebSocket-Protocol} header, and
            C{WebSocketHandler} instance. The I{WebSocket-Protocol} header will
            be C{None} if not specified by the client.
        """
        def finish():
            self.channel.transport.loseConnection()
        if self.queued:
            return finish()
        originHeaders = self.requestHeaders.getRawHeaders("Origin", [])
        if len(originHeaders) != 1:
            return finish()
        hostHeaders = self.requestHeaders.getRawHeaders("Host", [])
        if len(hostHeaders) != 1:
            return finish()

        handlerFactory = self.site.handlers.get(self.uri)
        if not handlerFactory:
            return finish()
        transport = WebSocketTransport(self)
        handler = handlerFactory(transport)
        transport._attachHandler(handler)

        protocolHeaders = self.requestHeaders.getRawHeaders(
            "WebSocket-Protocol", [])
        if len(protocolHeaders) not in (0,  1):
            return finish()
        if protocolHeaders:
            if protocolHeaders[0] not in self.site.supportedProtocols:
                return finish()
            protocolHeader = protocolHeaders[0]
        else:
            protocolHeader = None
        return originHeaders[0], hostHeaders[0], protocolHeader, handler

    def _getOneHeader(self, name):
        headers = self.requestHeaders.getRawHeaders(name)
        if not headers or len(headers) > 1:
            return None
        return headers[0]

    def _clientHandshakeHybi(self):
        """
        Initial handshake, as defined in hybi-10.

        If the client is not following the hybi-10 protocol or is requesting a
        version that's lower than what hybi-10 describes, the connection will
        be closed.

        Otherwise the appropriate transport and content decoders will be
        plugged in and the connection will be estabilished.
        """
        version = self._getOneHeader("Sec-WebSocket-Version")
        # we only speak version 8~13 of the protocol
        if int(version) > 13 or int(version) < 8:
            self.setResponseCode(426, "Upgrade Required")
            self.setHeader("Sec-WebSocket-Version", "8")
            return self.finish()

        key = self._getOneHeader("Sec-WebSocket-Key")
        if not key:
            self.setResponseCode(400, "Bad Request")
            return self.finish()

        handlerFactory = self.site.handlers.get(self.uri)
        if not handlerFactory:
            self.setResponseCode(404, "Not Found")
            return self.finish()

        transport = WebSocketHybiTransport(self)
        handler = handlerFactory(transport)
        transport._attachHandler(handler)

        accept = base64.b64encode(sha1(key + self.ACCEPT_GUID).digest())
        self.startedWriting = True
        handshake = [
            "HTTP/1.1 101 Switching Protocols",
            "Upgrade: websocket",
            "Connection: Upgrade",
            "Sec-WebSocket-Accept: %s" % accept]

        for header in handshake:
            self.write("%s\r\n" % header)

        self.write("\r\n")
        self.channel.setRawMode()
        self.channel._transferDecoder = WebSocketHybiFrameDecoder(
            self, handler)
        handler.transport._connectionMade()

    def renderWebSocket(self):
        """
        Render a WebSocket request.

        If the request is not identified with a proper WebSocket handshake, the
        connection will be closed. Otherwise, the response to the handshake is
        sent and a C{WebSocketHandler} is created to handle the request.
        """
        # check for hybi handshake requests
        if self.requestHeaders.hasHeader("Sec-WebSocket-Version"):
            return self._clientHandshakeHybi()

        # check for post-75 handshake requests
        isSecHandshake = self.requestHeaders.getRawHeaders("Sec-WebSocket-Key1", [])
        if isSecHandshake:
            self._clientHandshake76()
        else:
            check = self._checkClientHandshake()
            if check is None:
                return
            originHeader, hostHeader, protocolHeader, handler = check
            self.startedWriting = True
            handshake = [
                "HTTP/1.1 101 Web Socket Protocol Handshake",
                "Upgrade: WebSocket",
                "Connection: Upgrade"]
            handshake.append("WebSocket-Origin: %s" % (originHeader))
            if self.isSecure():
                scheme = "wss"
            else:
                scheme = "ws"
            handshake.append(
                "WebSocket-Location: %s://%s%s" % (
                scheme, hostHeader, self.uri))

            if protocolHeader is not None:
                handshake.append("WebSocket-Protocol: %s" % protocolHeader)

            for header in handshake:
                self.write("%s\r\n" % header)

            self.write("\r\n")
            self.channel.setRawMode()
            # XXX we probably don't want to set _transferDecoder
            self.channel._transferDecoder = WebSocketFrameDecoder(
                self, handler)
            handler.transport._connectionMade()
            return



class WebSocketSite(Site):
    """
    @ivar handlers: a C{dict} of names to L{WebSocketHandler} factories.
    @type handlers: C{dict}
    @ivar supportedProtocols: a C{list} of supported I{WebSocket-Protocol}
        values. If a value is passed at handshake and doesn't figure in this
        list, the connection is closed.
    @type supportedProtocols: C{list}
    """
    requestFactory = WebSocketRequest

    def __init__(self, resource, logPath=None, timeout=60*60*12,
                 supportedProtocols=None):
        Site.__init__(self, resource, logPath, timeout)
        self.handlers = {}
        self.supportedProtocols = supportedProtocols or []

    def addHandler(self, name, handlerFactory):
        """
        Add or override a handler for the given C{name}.

        @param name: the resource name to be handled.
        @type name: C{str}
        @param handlerFactory: a C{WebSocketHandler} factory.
        @type handlerFactory: C{callable}
        """
        if not name.startswith("/"):
            raise ValueError("Invalid resource name.")
        self.handlers[name] = handlerFactory



class WebSocketTransport(object):
    """
    Transport abstraction over WebSocket, providing classic Twisted methods and
    callbacks.
    """
    implements(interfaces.ITransport, interfaces.IPushProducer)

    _handler = None

    def __init__(self, request):
        self._request = request
        self._connected = 1
        self._paused = False
        self._pending_frames = []
        self._request.notifyFinish().addErrback(self._connectionLost)

    def _attachHandler(self, handler):
        """
        Attach the given L{WebSocketHandler} to this transport.
        """
        self._handler = handler

    def _connectionMade(self):
        """
        Called when a connection is made.
        """
        self._handler.connectionMade()

    def _connectionLost(self, reason):
        """
        Forward connection lost event to the L{WebSocketHandler}.
        """
        self._connected = 0
        self._paused = False
        self._pending_frames = []
        self._handler.connectionLost(reason)
        del self._request.transport
        del self._request
        del self._handler

    def getPeer(self):
        """
        Return a tuple describing the other side of the connection.

        @rtype: C{tuple}
        """
        return self._request.transport.getPeer()

    def getHost(self):
        """
        Similar to getPeer, but returns an address describing this side of the
        connection.

        @return: An L{IAddress} provider.
        """
        return self._request.transport.getHost()

    def write(self, frame):
        """
        Send the given frame to the connected client.

        @param frame: a I{UTF-8} encoded C{str} to send to the client.
        @type frame: C{str}
        """
        if self._connected:
            if self._paused:
                self._pending_frames.append(frame)
            else:
                self._request.write("\x00%s\xff" % frame)

    def writeSequence(self, frames):
        """
        Send a sequence of frames to the connected client.
        """
        if self._connected:
            self._request.write("".join(["\x00%s\xff" % f for f in frames]))

    def loseConnection(self):
        """
        Close the connection.
        """
        if self._connected:
            self._request.transport.loseConnection()

    def pauseProducing(self):
        self._paused = True
        self._handler.pauseProducing()

    def resumeProducing(self):
        self._paused = False
        self._sendPendingFrames()
        self._handler.resumeProducing()

    def _sendPendingFrames(self):
        """
        Write pending frames
        """
        if not self._paused:
            self.writeSequence(self._pending_frames)
            self._pending_frames = []

class WebSocketHybiTransport(WebSocketTransport):
    """
    A WebSocket transport that speaks the hybi-10 protocol. The L{ITransport}
    methods are set up to send Text frames containing the payload. To have
    finer-grained control over the type of frame being sent, the transport
    provides a L{sendFrame} method.
    """
    def write(self, frame):
        """
        Treat the given frame as a text frame and send it to the client.

        @param frame: a I{UTF-8} encoded C{str} to send to the client.
        @type frame: C{str}
        """
        self.sendFrame(OPCODE_TEXT, frame)

    def writeSequence(self, frames):
        """
        Send a sequence of text frames to the connected client.
        """
        for frame in frames:
            self.sendFrame(OPCODE_TEXT, frame)

    def sendFrame(self, opcode, payload, fragmented=False):
        """
        Send a frame with the given opcode and payload to the client. If the
        L{fragmented} parameter is set, the message frame will contain a flag
        saying it's part of a fragmented payload, by default data is sent as a
        self-contained frame. Note that if you use fragmentation support, it is
        up to you to correctly set the first frame's opcode and then use
        L{OPCODE_CONT} on the following continuation frames.

        Payloads sent using this method are never masked.

        @param opcode: the opcode as defined in hybi-10
        @type opcode: C{int}
        @param payload: the frame's payload
        @type payload: C{str}
        @param fragmented: should the frame be marked as part of a fragmented payload
        @type fragmented: C{bool}
        """
        if opcode not in ALL_OPCODES:
            raise ValueError("Invalid opcode 0x%X" % opcode)

        length = len(payload)

        # there's always the header and at least one length field
        spec = ">BB"
        if fragmented:
            header = 0x00
        else:
            header = 0x80
        data = [header | opcode]

        # there's no masking, so the high bit of the first byte of length is
        # always 0
        if 125 < length <= 65535:
            # add a 16-bit int to the spec and append 126 value, which means
            # "interpret the next two bytes"
            spec += "H"
            data.append(126)
        elif length > 65535:
            # same for even longer frames
            spec += "Q"
            data.append(127)

        data.append(length)
        header = struct.pack(spec, *data)
        self._request.write(header + payload)


class WebSocketHandler(object):
    """
    Base class for handling WebSocket connections. It mainly provides a
    transport to send frames, and a callback called when frame are received,
    C{frameReceived}.

    @ivar transport: a C{WebSocketTransport} instance.
    @type: L{WebSocketTransport}
    """

    implements(interfaces.IPushProducer)

    def __init__(self, transport):
        """
        Create the handler, with the given transport
        """
        self.transport = transport
        self._paused = False
        self._pending_frames = []

    def _frameReceived(self, frame):
        if self._paused:
            self._pending_frames.append(frame)
        else:
            self.frameReceived(frame)

    def _processPedingFrames(self):
        if not self._paused:
            for frame in self._pending_frames:
                self.frameReceived(frame)
            self._pending_frames = []

    def pauseProducing(self):
        self._paused = 1

    def resumeProducing(self):
        self._paused = 0
        self._processPedingFrames()

    def frameReceived(self, frame):
        """
        Called when a frame is received.

        @param frame: a I{UTF-8} encoded C{str} sent by the client.
        @type frame: C{str}
        """


    def binaryFrameReceived(self, data):
        """
        Called when a binary is received via the hybi protocol.

        @param data: a binary C{str} sent by the client.
        @type data: C{str}
        """


    def pongReceived(self, data):
        """
        Called when a pong control message is received via the hybi protocol.

        @param data: the payload sent by the client.
        @type data: C{str}
        """


    def closeReceived(self, code, msg):
        """
        Called when a close control message is received via the hybi protocol.

        @param code: the status code of the close message, if present
        @type code: C{int} or C{None}
        @param msg: the I{UTF-8} encoded message sent by the client, if present
        @type msg: C{str} or C{None}
        """


    def frameLengthExceeded(self):
        """
        Called when too big a frame is received. The default behavior is to
        close the connection, but it can be customized to do something else.
        """
        self.transport.loseConnection()


    def connectionMade(self):
        """
        Called when a connection is made.
        """

    def connectionLost(self, reason):
        """
        Callback called when the underlying transport has detected that the
        connection is closed.
        """


class IncompleteFrame(Exception):
    """
    Not enough data to complete a WebSocket frame.
    """


class DecodingError(Exception):
    """
    The incoming data is not valid WebSocket protocol data.
    """


class WebSocketFrameDecoder(object):
    """
    Decode WebSocket frames and pass them to the attached C{WebSocketHandler}
    instance.

    @ivar MAX_LENGTH: maximum len of a text frame allowed, before calling
        C{frameLengthExceeded} on the handler.
    @type MAX_LENGTH: C{int}
    @ivar MAX_BINARY_LENGTH: like C{MAX_LENGTH}, but for 0xff type frames
    @type MAX_BINARY_LENGTH: C{int}
    @ivar closing: a flag set when the closing handshake has been received
    @type closing: C{bool}
    @ivar request: C{Request} instance.
    @type request: L{twisted.web.server.Request}
    @ivar handler: L{WebSocketHandler} instance handling the request.
    @type handler: L{WebSocketHandler}
    @ivar _data: C{list} of C{str} buffering the received data.
    @type _data: C{list} of C{str}
    @ivar _currentFrameLength: length of the current handled frame, plus the
        additional leading byte.
    @type _currentFrameLength: C{int}
    """

    MAX_LENGTH = 16384
    MAX_BINARY_LENGTH = 2147483648
    closing = False

    def __init__(self, request, handler):
        self.request = request
        self.handler = handler
        self.closing = False
        self._data = []
        self._currentFrameLength = 0
        self._state = "FRAME_START"

    def dataReceived(self, data):
        """
        Parse data to read WebSocket frames.

        @param data: data received over the WebSocket connection.
        @type data: C{str}
        """
        if not data or self.closing:
            return
        self._data.append(data)

        while self._data and not self.closing:
            try:
                self.consumeData(self._data[-1])
            except IncompleteFrame:
                break
            except DecodingError:
                log.err()
                self.request.transport.loseConnection()
                break

    def consumeData(self, data):
        """
        Process the last data chunk received.

        After processing is done, L{IncompleteFrame} should be raised or
        L{_addRemainingData} should be called.

        @param data: last chunk of data received.
        @type data: C{str}
        """
    consumeData = makeStatefulDispatcher("consumeData", consumeData)

    def _consumeData_FRAME_START(self, data):
        self._currentFrameLength = 0

        if data[0] == "\x00":
            self._state = "PARSING_TEXT_FRAME"
        elif data[0] == "\xff":
            self._state = "PARSING_LENGTH"
        else:
            raise DecodingError("Invalid frame type 0x%s" %
                                data[0].encode("hex"))

        self._addRemainingData(data[1:])

    def _consumeData_PARSING_TEXT_FRAME(self, data):
        endIndex = data.find("\xff")
        if endIndex == -1:
            self._currentFrameLength += len(data)
        else:
            self._currentFrameLength += endIndex

        self._currentFrameLength += endIndex
        # check length + 1 to account for the initial frame type byte
        if self._currentFrameLength + 1 > self.MAX_LENGTH:
            self.handler.frameLengthExceeded()

        if endIndex == -1:
            raise IncompleteFrame()

        frame = "".join(self._data[:-1]) + data[:endIndex]
        self.handler._frameReceived(frame)

        remainingData = data[endIndex + 1:]
        self._addRemainingData(remainingData)

        self._state = "FRAME_START"

    def _consumeData_PARSING_LENGTH(self, data):
        current = 0
        available = len(data)

        while current < available:
            byte = ord(data[current])
            length, more = byte & 0x7F, bool(byte & 0x80)

            if not length:
                self._closingHandshake()
                raise IncompleteFrame()

            self._currentFrameLength *= 128
            self._currentFrameLength += length

            current += 1

            if not more:
                if self._currentFrameLength > self.MAX_BINARY_LENGTH:
                    self.handler.frameLengthExceeded()

                remainingData = data[current:]
                self._addRemainingData(remainingData)
                self._state = "PARSING_BINARY_FRAME"
                break
        else:
            raise IncompleteFrame()

    def _consumeData_PARSING_BINARY_FRAME(self, data):
        available = len(data)

        if self._currentFrameLength <= available:
            remainingData = data[self._currentFrameLength:]
            self._addRemainingData(remainingData)
            self._state = "FRAME_START"
        else:
            self._currentFrameLength -= available
            self._data[:] = []

    def _addRemainingData(self, remainingData):
        if remainingData:
            self._data[:] = [remainingData]
        else:
            self._data[:] = []

    def _closingHandshake(self):
        self.closing = True
        # send the closing handshake
        self.request.transport.write("\xff\x00")
        # discard all buffered data
        self._data[:] = []


class WebSocketHybiFrameDecoder(WebSocketFrameDecoder):

    def __init__(self, request, handler):
        WebSocketFrameDecoder.__init__(self, request, handler)
        self._opcode = None
        self._fragment_opcode = None
        self._fragments = []
        self._state = "HYBI_FRAME_START"

    def _consumeData_HYBI_FRAME_START(self, data):
        self._opcode = None

        byte = ord(data[0])
        fin, reserved, opcode = byte & 0x80, byte & 0x70, byte & 0x0F

        if reserved:
            raise DecodingError("Reserved bits set: 0x%02X" % byte)

        if opcode not in ALL_OPCODES:
            raise DecodingError("Invalid opcode 0x%X" % opcode)

        if not fin:
            # part of a fragmented frame
            if not self._fragment_opcode:
                # first of the fragmented frames, which determines the opcode
                if opcode not in DATA_OPCODES:
                    raise DecodingError(
                        "Fragmented frame with invalid opcode 0x%X" % opcode)
                # save the opcode for later use
                self._fragment_opcode = opcode
            else:
                # already reading a fragmet, and this is a fragmented frame, so
                # it has to use the continuation opcode
                if opcode != OPCODE_CONT:
                    raise DecodingError(
                        "Continuation frame with invalid opcode 0x%X" % opcode)
        else:
            # self-contained frame or last of the fragmented frames
            if self._fragment_opcode:
                # a fragmented frame is pending, so this can only be the end of
                # it or a control message
                if opcode not in CONTROL_OPCODES and opcode != OPCODE_CONT:
                    raise DecodingError(
                        "Final frame with invalid opcode 0x%X" % opcode)
            else:
                # no fragmented frames pending, so this cannot be a
                # continuation frame
                if opcode == OPCODE_CONT:
                    raise DecodingError(
                        "Final frame with invalid opcode 0x%X" % opcode)
            self._opcode = opcode

        self._state = "HYBI_PARSING_LENGTH"
        self._addRemainingData(data[1:])

    def _consumeData_HYBI_PARSING_LENGTH(self, data):
        byte = ord(data[0])
        masked, length = byte & 0x80, byte & 0x7F

        if not masked:
            raise DecodingError("Unmasked frame received")

        if length < 126:
            self._currentFrameLength = length
            self._state = "HYBI_MASKING_KEY"
        elif length == 126:
            self._state = "HYBI_PARSING_LENGTH_2"
        elif length == 127:
            self._state = "HYBI_PARSING_LENGTH_3"

        self._addRemainingData(data[1:])

    def _consumeData_HYBI_PARSING_LENGTH_2(self, data):
        self._parse_length_spec(2, ">H")

    def _consumeData_HYBI_PARSING_LENGTH_3(self, data):
        self._parse_length_spec(8, ">Q", 0x7fffffffffffffff)

    def _parse_length_spec(self, needed, spec, limit=None):
        # if the accumulated data is not long enough to parse out the length,
        # keep on accumulating
        if sum(map(len, self._data)) < needed:
            raise IncompleteFrame()

        data = "".join(self._data)
        self._currentFrameLength = struct.unpack(spec, data[:needed])[0]
        if limit and self._currentFrameLength > limit:
            raise DecodingError(
                "Frame length exceeded: %r" % self._currentFrameLength)
        self._addRemainingData(data[needed:])

        self._state = "HYBI_MASKING_KEY"

    def _consumeData_HYBI_MASKING_KEY(self, data):
        if sum(map(len, self._data)) < 4:
            raise IncompleteFrame()

        data = "".join(self._data)
        self._maskingKey = struct.unpack(">4B", data[:4])
        self._addRemainingData(data[4:])

        if self._currentFrameLength:
            self._state = "HYBI_PAYLOAD"
        else:
            # there will be no payload, notify the handler of an empty frame
            # and continue
            self._frameCompleted("", data[4:])

    def _consumeData_HYBI_PAYLOAD(self, data):
        available = len(data)

        if self._currentFrameLength > available:
            self._currentFrameLength -= available
            raise IncompleteFrame()

        frame = "".join(self._data[:-1]) + data[:self._currentFrameLength]

        # unmask the frame
        bufferedPayload = itertools.chain(*self._data[:-1])
        restOfPayload = data[:self._currentFrameLength]
        allData = itertools.chain(bufferedPayload, restOfPayload)

        key = itertools.cycle(self._maskingKey)

        def xor(c, k):
            return chr(ord(c) ^ k)
        unmasked = itertools.imap(xor, allData, key)

        frame = "".join(unmasked)
        remainingData = data[self._currentFrameLength:]

        self._frameCompleted(frame, remainingData)

    def _frameCompleted(self, frame, remainingData):
        # if it's part of a fragmented frame, store the payload
        if self._opcode is None:
            self._fragments.append(frame)

        # if it's the last of the fragmented frames, replace the opcode with
        # the original one from the fragment and the frame with the accumulated
        # payload
        if self._opcode == OPCODE_CONT:
            self._opcode = self._fragment_opcode
            self._fragments.append(frame)
            frame = "".join(self._fragments)
            self._fragment_opcode = None
            self._fragments[:] = []

        if self._opcode == OPCODE_TEXT:
            # assume it's valid UTF-8 and let the client handle the rest
            if len(frame) > self.MAX_LENGTH:
                self.handler.frameLengthExceeded()
            self.handler._frameReceived(frame)
        elif self._opcode == OPCODE_BINARY:
            if len(frame) > self.MAX_BINARY_LENGTH:
                self.handler.frameLengthExceeded()
            self.handler.binaryFrameReceived(frame)
        elif self._opcode == OPCODE_PING:
            self.handler.transport.sendFrame(OPCODE_PONG, frame)
        elif self._opcode == OPCODE_PONG:
            self.handler.pongReceived(frame)

        self._state = "HYBI_FRAME_START"
        self._addRemainingData(remainingData)

        # if the opcode was CLOSE, initiate connection closing
        if self._opcode == OPCODE_CLOSE:
            self._hybiClose(frame)

    def _hybiClose(self, frame):
        self.closing = True

        # try to parse out the status code and message
        if len(frame) > 1:
            code = struct.unpack(">H", frame[:2])[0]
            msg = frame[2:]
        else:
            code, msg = None, None
        # let the handler know
        self.handler.closeReceived(code, msg)

        # send the closing handshake
        self.handler.transport.sendFrame(OPCODE_CLOSE, "")

        # discard all buffered data and lose connection
        self._data[:] = []
        self.handler.transport.loseConnection()

class WebSocketWrapperProtocol(ProtocolWrapper):
    pass

class WebSocketWrapperFactory(WrappingFactory):
    """
    Factory which wraps another factory to provide WebSockets transports for
    all of its protocols.
    """

    protocol = ProtocolWrapper

    def buildHandler(self, transport):
        """"
        Called in order to build C{WebSocketHandler} objects.
    
        @ivar transport: a C{WebSocketTransport} instance.
        @type: L{WebSocketTransport}
        """
        
        protocol = self.buildProtocol(transport.getHost())

        handler = WebSocketHandler(transport)

        def connectionMade():
            protocol.makeConnection(transport)
            
        def frameReceived(frame):
            protocol.dataReceived(frame)

        def connectionLost(reason):
            protocol.connectionLost(reason)
        
        handler.connectionMade = connectionMade
        handler.frameReceived = frameReceived
        handler.connectionLost = connectionLost
        
        return handler

__all__ = ["WebSocketHandler", "WebSocketSite", "WebSocketWrapperFactory", "WebSocketWrapperProtocol"]