def setUp(self): self.channel = DummyChannel() request = Request(self.channel, False) transport = WebSocketTransport(request) handler = TestHandler(transport) transport._attachHandler(handler) self.decoder = WebSocketFrameDecoder(request, handler) # taken straight from the IETF spec, masking added where appropriate self.hello = "\x81\x85\x37\xfa\x21\x3d\x7f\x9f\x4d\x51\x58" self.frag_hello = ("\x01\x83\x12\x21\x65\x23\x5a\x44\x09", "\x80\x82\x63\x34\xf1\x00\x0f\x5b") self.binary_orig = "\x3f" * 256 self.binary = ("\x82\xfe\x01\x00\x12\x6d\xa6\x23" + "\x2d\x52\x99\x1c" * 64) self.long_binary = ("\x82\xff\x00\x00\x00\x00\x00\x01\x00\x00" + "\x12\x6d\xa6\x23" + "\x2d\x52\x99\x1c" * 16384) self.long_binary_orig = "\x3f" * 65536 self.ping = "\x89\x85\x56\x23\x88\x23\x1e\x46\xe4\x4f\x39" self.pong = "\x8a\x85\xde\x41\x0f\x34\x96\x24\x63\x58\xb1" self.pong_unmasked = "\x8a\x05\x48\x65\x6c\x6c\x6f" # code 1000, message "Normal Closure" self.close = ("\x88\x90\x34\x23\x87\xde\x37\xcb\xc9\xb1\x46" "\x4e\xe6\xb2\x14\x60\xeb\xb1\x47\x56\xf5\xbb") ## close message can be empty or with normal close code (1000) self.empty_unmasked_close_list = ("\x88\x00", "\x88\x02\x03\xe8") self.empty_text = "\x81\x80\x00\x01\x02\x03" self.cont_empty_text = "\x00\x80\x00\x01\x02\x03"
class WebSocketFrameDecoderXor(TestCase): """ Test the vitamined xor function of the decoder (_xxor) """ def setUp(self): self.channel = DummyChannel() request = Request(self.channel, False) transport = WebSocketTransport(request) handler = WebSocketHandler(transport) self.decoder = WebSocketFrameDecoder(request, handler) self.decoder._mask = "Key!" self.decoder._rkeys = {} self.decoder._maskIndex = 0 def xor(self, data): """ Reference version """ return ''.join(chr(ord(x) ^ ord(y)) for (x,y) in izip(data, cycle(self.decoder._mask))) def test_xxorShort(self): """ Short data (less than 16 bytes) """ data = "Hello" self.assertEquals(self.decoder._xxor(data), self.xor(data)) def test_xxorLong(self): """ Long data (more than 16 bytes). Make sure the data is not a multiple of 8 Test the long xor capability """ data = "Hello" * 500 + "trail" self.assertEquals(self.decoder._xxor(data), self.xor(data)) def test_xxorSplit(self): """ Long data (4005 bytes), processed in chunks of 37 bytes Test the rolling key mechanism """ data = "Hello" * 500 + "trail" chunk_size = 37 self.assertEquals("".join(self.decoder._xxor(data[i:i+chunk_size]) for i in range(0, len(data), chunk_size)), self.xor(data))
def setUp(self): self.channel = DummyChannel() request = Request(self.channel, False) transport = WebSocketTransport(request) handler = WebSocketHandler(transport) self.decoder = WebSocketFrameDecoder(request, handler) self.decoder._mask = "Key!" self.decoder._rkeys = {} self.decoder._maskIndex = 0
def connectionMade(self): self.data = "" self.extensions = [] self.isSecure = (self.factory.scheme == "wss") self.transport.connectionLost = self.connectionLost self.sendHeader() self.decoder = WebSocketFrameDecoder(self.transport, self, isClient=True, assembleFragments=self.assembleFragments) self._firstFragment = True
class WebSocketServerFrameDecoderTestCase(TestCase): """ Test for C{WebSocketFrameDecoder}. """ def setUp(self): self.channel = DummyChannel() request = Request(self.channel, False) transport = WebSocketTransport(request) handler = TestHandler(transport) transport._attachHandler(handler) self.decoder = WebSocketFrameDecoder(request, handler) # taken straight from the IETF spec, masking added where appropriate self.hello = "\x81\x85\x37\xfa\x21\x3d\x7f\x9f\x4d\x51\x58" self.frag_hello = ("\x01\x83\x12\x21\x65\x23\x5a\x44\x09", "\x80\x82\x63\x34\xf1\x00\x0f\x5b") self.binary_orig = "\x3f" * 256 self.binary = ("\x82\xfe\x01\x00\x12\x6d\xa6\x23" + "\x2d\x52\x99\x1c" * 64) self.long_binary = ("\x82\xff\x00\x00\x00\x00\x00\x01\x00\x00" + "\x12\x6d\xa6\x23" + "\x2d\x52\x99\x1c" * 16384) self.long_binary_orig = "\x3f" * 65536 self.ping = "\x89\x85\x56\x23\x88\x23\x1e\x46\xe4\x4f\x39" self.pong = "\x8a\x85\xde\x41\x0f\x34\x96\x24\x63\x58\xb1" self.pong_unmasked = "\x8a\x05\x48\x65\x6c\x6c\x6f" # code 1000, message "Normal Closure" self.close = ("\x88\x90\x34\x23\x87\xde\x37\xcb\xc9\xb1\x46" "\x4e\xe6\xb2\x14\x60\xeb\xb1\x47\x56\xf5\xbb") ## close message can be empty or with normal close code (1000) self.empty_unmasked_close_list = ("\x88\x00", "\x88\x02\x03\xe8") self.empty_text = "\x81\x80\x00\x01\x02\x03" self.cont_empty_text = "\x00\x80\x00\x01\x02\x03" def assertOneDecodingError(self): """ Assert that exactly one L{DecodingError} has been logged and return that error. """ errors = self.flushLoggedErrors(DecodingError) self.assertEquals(len(errors), 1) return errors[0] def test_oneTextFrame(self): """ We can send one frame handled with one C{dataReceived} call. """ self.decoder.dataReceived(self.hello) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_chunkedTextFrame(self): """ We can send one text frame handled with multiple C{dataReceived} calls. """ # taken straight from the IETF spec for part in (self.hello[:1], self.hello[1:3], self.hello[3:7], self.hello[7:]): self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_fragmentedTextFrame(self): """ We can send a fragmented frame handled with one C{dataReceived} call. """ self.decoder.dataReceived("".join(self.frag_hello)) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_chunkedfragmentedTextFrame(self): """ We can send a fragmented text frame handled with multiple C{dataReceived} calls. """ # taken straight from the IETF spec for part in (self.frag_hello[0][:3], self.frag_hello[0][3:]): self.decoder.dataReceived(part) for part in (self.frag_hello[1][:1], self.frag_hello[1][1:]): self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_twoFrames(self): """ We can send two frames together and they will be correctly parsed. """ self.decoder.dataReceived("".join(self.frag_hello) + self.hello) self.assertEquals(self.decoder.handler.frames, ["Hello"] * 2) def test_controlInterleaved(self): """ A control message (in this case a pong) can appear between the fragmented frames. """ data = self.frag_hello[0] + self.pong + self.frag_hello[1] for part in data[:2], data[2:7], data[7:8], data[8:14], data[14:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.frames, ["Hello"]) self.assertEquals(self.decoder.handler.pongs, ["Hello"]) def test_binaryFrameAllAtOnce(self): """ We can send a binary frame that uses a longer length field. """ self.decoder.dataReceived(self.binary) self.assertEquals(self.decoder.handler.binaryFrames, [self.binary_orig]) def test_binaryFrame(self): """ We can send a binary frame that uses a longer length field. Split data in some interesting locations. """ data = self.binary + self.hello for part in data[:3], data[3:4], data[4:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.binaryFrames, [self.binary_orig]) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_binaryFrameMultiSplit(self): """ We can send a binary frame that uses a longer length field, plus one other frame. Split in all header locations. """ data = self.binary + self.hello for i in range(10): self.decoder.handler.binaryFrames = [] self.decoder.handler.frames = [] for part in data[:i], data[i:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.binaryFrames, [self.binary_orig], "Failed at i=%d" % i) self.assertEquals(self.decoder.handler.frames, ["Hello"]) for i in range(10): self.decoder.handler.binaryFrames = [] self.decoder.handler.frames = [] for part in data[:len(self.binary)+i], data[len(self.binary)+i:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.binaryFrames, [self.binary_orig], "Failed (last) at i=%d" % i) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_binaryFrameOneByOne(self): """ We can send a binary frame that uses a longer length field. Bytes are sent one by one. """ data = self.binary + self.hello for part in data: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.binaryFrames, [self.binary_orig]) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_longBinaryFrame(self): """ We can send a binary frame that uses a very long length field. """ data = self.long_binary + self.hello for part in data[:3], data[3:4], data[4:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.binaryFrames, [self.long_binary_orig]) self.assertEquals(self.decoder.handler.frames, ["Hello"]) def test_pingInterleaved(self): """ We can get a ping frame in the middle of a fragmented frame and we'll correctly send a pong resonse. """ data = self.frag_hello[0] + self.ping + self.frag_hello[1] for part in data[:12], data[12:16], data[16:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.frames, ["Hello"]) result = self.channel.transport.written.getvalue() headers, response = result.split('\r\n\r\n') self.assertEquals(response, self.pong_unmasked) def test_pingInterleavedOneByOne(self): """ We can get a ping frame in the middle of a fragmented frame and we'll correctly send a pong resonse. Send bytes one by one. """ data = self.frag_hello[0] + self.ping + self.frag_hello[1] for part in data: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.frames, ["Hello"]) result = self.channel.transport.written.getvalue() headers, response = result.split('\r\n\r\n') self.assertEquals(response, self.pong_unmasked) def test_close(self): """ A close frame causes the remaining data to be discarded and the connection to be closed. """ self.decoder.dataReceived(self.hello + self.close + "crap" * 20) self.assertEquals(self.decoder.handler.frames, ["Hello"]) self.assertEquals(self.decoder.handler.closes, [(1000, "Normal Closure")]) result = self.channel.transport.written.getvalue() headers, response = result.split('\r\n\r\n') self.assertIn(response, self.empty_unmasked_close_list) self.assertTrue(self.channel.transport.disconnected) def test_emptyFrame(self): """ An empty text frame is correctly parsed. """ self.decoder.dataReceived(self.empty_text) self.assertEquals(self.decoder.handler.frames, [""]) def test_emptyFrameInterleaved(self): """ An empty fragmented frame and a interleaved pong message are received and parsed. """ data = (self.frag_hello[0] + self.cont_empty_text + self.pong + self.frag_hello[1]) for part in data[:1], data[1:8], data[8:17], data[17:]: self.decoder.dataReceived(part) self.assertEquals(self.decoder.handler.frames, ["Hello"]) self.assertEquals(self.decoder.handler.pongs, ["Hello"])
class WebSocketClient(Protocol, TimeoutMixin, _PauseableMixin): """ Protocol for websocket clients. @ivar connectionEstablished: a callable which will be invoqued when the connection has been established. @ivar handshakeError: A one-argument callable which will be invoked when the handshake failed, the argument being a string describing the reason for the failure. @ivar frameReceived: A one-argument callable which will be invoked when the terminal chunk is received. It will be invoked with all bytes which were delivered to this protocol which came after the terminal chunk. @ivar fragmentReceived: A one-argument callable which will be invoked each time a fragment of data is received for streaming purpose. The argument is the plain data as received on the wire. @ivar dataError: A one-argument callable which will be invoked when the handshake failed, the argument being a string describing the reason for the failure. """ extensions = [] assembleFragments = False def connectionMade(self): self.data = "" self.extensions = [] self.isSecure = (self.factory.scheme == "wss") self.transport.connectionLost = self.connectionLost self.sendHeader() self.decoder = WebSocketFrameDecoder(self.transport, self, isClient=True, assembleFragments=self.assembleFragments) self._firstFragment = True def sendHeader(self): self.in_handshake = True self.fields = [ "GET %s HTTP/1.1\r\n" % self.factory.path, "Host: %s:%s\r\n" % (self.factory.host, self.factory.port), ## XXX: Can be added later: Proxy-authorization ] if self.isSecure: if not has_ssl: self.handshakeError("The PyOpenSSL library is not installed." " Cannot use the wss scheme") self.transport.loseConnection() return self.transport.startTLS(ClientTLSContext(), self.factory) else: self.sendHeaderEnd() def sendHeaderEnd(self): fields = self.fields self.nonce = b64encode(generate_key()) ## These fields must be shuffled field_list = [ "Upgrade: WebSocket\r\n", "Connection: Upgrade\r\n", "Sec-WebSocket-Key: %s\r\n" % self.nonce, "Origin: http://%s\r\n" % self.factory.origin, "Sec-WebSocket-Version: 13\r\n" ] if self.factory.subprotocolsAvailable: field_list.append("Sec-WebSocket-Protocol: %s\r\n" % ",".join(self.factory.subprotocolsAvailable)) if self.factory.extensionsAvailable: field_list.append("Sec-WebSocket-Extensions: %s\r\n" % ",".join(ext.requestHeader for ext in self.factory.extensionsAvailable)) if self.factory.extra_headers != None: field_list.extend(self.factory.extra_headers) field_list.append("\r\n") all_fields = ''.join(fields + field_list) self.transport.write(all_fields) def writeBinary(self, data, fragmented=False): """ Treat the given frame as a binary frame and send it to the client. @param frame: a binary C{str} to send to the client. @type frame: C{str} """ self.sendFrame(OPCODE_BINARY, data, fragmented) def write(self, data, fragmented=False): """ Treat the given utf-8 data 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, data.encode('utf-8'), fragmented=fragmented) def writeFragment(self, data, final=False): """ Send a fragment of a text frame to the client. @param fragment: a I{UTF-8} encoded C{str} to send to the client. @type fragment: C{str} """ self.sendFrame(OPCODE_TEXT, data.encode('utf-8'), fragmented=True, final=final) write.writeFragment = writeFragment def writeBinaryFragment(self, data, final=False): """ Send a fragment of a binary frame to the client. @param fragment: a C{str} to send to the client. @type fragment: C{str} """ self.sendFrame(OPCODE_BINARY, data, fragmented=True, final=final) writeBinary.writeFragment = writeBinaryFragment def sendPing(self, data=""): """ Send a simple ping packet. The factory should implement pongReceived to check the pong response """ self.sendFrame(OPCODE_PING, data) def writeSequence(self, frames): """ Send a sequence of text frames to the connected client. """ for frame in frames: self.sendFrame(OPCODE_TEXT, frame) def writeBinarySequence(self, frames): """ Send a sequence of binary frames to the connected client. """ for frame in frames: self.sendFrame(OPCODE_BINARY, frame) def sendFrame(self, opcode, payload, fragmented=False, final=False, **kwargs): """ 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 rfc6455 @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 ## Process extensions if self.extensions: if opcode == OPCODE_TEXT: for ext in self.extensions: payload = ext.processOutgoingFrame(payload, **kwargs) elif opcode == OPCODE_BINARY: for ext in self.extensions: payload = ext.processOutgoingBinaryFrame(payload, **kwargs) length = len(payload) if fragmented: if final: if self._firstFragment: ## Special case when final fragment is also ## the first one: same as unfragmented data = [0x80 | opcode] else: ## Final fragment, opcode set to 0 data = [0x80] self._firstFragment = True elif self._firstFragment: ## First fragment, final set to 0 data = [opcode] self._firstFragment = False else: ## Continuation data = [OPCODE_CONT] else: data = [0x80 | opcode] ## Create the frame header, depending on the length if length < 126: spec = "!BBI" data.append(0x80|length) elif length > 65535: # same for even longer frames spec = "!BBQI" data.append(0xff) data.append(length) else: ## add a 16-bit int to the spec and append 126 value, which means ## "interpret the next two bytes" spec = "!BBHI" data.append(0xfe) data.append(length) ## Add the mask mask = getrandbits(32) data.append(mask) header = pack(spec, *data) payload = header + xxor(mask, payload) self.transport.write(payload) def dataReceived(self, data): if self.in_handshake: self.data += data try: remainingData = self.checkHandshake() except HandshakeResponseError, error: self.handshakeError(error) self.transport.loseConnection() return ## Header is over, witch to data handling mode self.in_handshake = False self.data = "" ## Just got response from server, might send something back self.connectionEstablished() self.decoder.dataReceived(remainingData) else: