def test_simple(self): client = self._setup_connection() frame = websockets.Frame.from_file(client.rfile) assert frame.payload == b'server-foobar' client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'client-foobar'))) client.wfile.flush() frame = websockets.Frame.from_file(client.rfile) assert frame.payload == b'client-foobar' client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef'))) client.wfile.flush() frame = websockets.Frame.from_file(client.rfile) assert frame.payload == b'\xde\xad\xbe\xef' client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) client.wfile.flush() assert len(self.master.state.flows) == 2 assert isinstance(self.master.state.flows[0], HTTPFlow) assert isinstance(self.master.state.flows[1], WebSocketFlow) assert len(self.master.state.flows[1].messages) == 5 assert self.master.state.flows[1].messages[0].content == b'server-foobar' assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT assert self.master.state.flows[1].messages[1].content == b'client-foobar' assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT assert self.master.state.flows[1].messages[2].content == b'client-foobar' assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef' assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef' assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
def test_change_payload(self): class Addon: def websocket_message(self, f): f.messages[-1].content = "foo" self.proxy.set_addons(Addon()) self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'foo' self.client.wfile.write( bytes( websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'foo' self.client.wfile.write( bytes( websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'foo'
def test_pong(self): self.setup_connection() self.client.wfile.write( bytes( websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.PING, payload=b'foobar'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) websockets.Frame.from_file(self.client.rfile) self.client.wfile.write( bytes( websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.flush() assert frame.header.opcode == websockets.OPCODE.PONG assert frame.payload == b'foobar' for i in range(20): if self.master.has_log("Pong Received from server", "info"): break time.sleep(0.01) else: raise AssertionError("No pong seen")
def handle_websockets(cls, rfile, wfile): wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'server-foobar'))) wfile.flush() frame = websockets.Frame.from_file(rfile) wfile.write(bytes(websockets.Frame(fin=1, opcode=frame.header.opcode, payload=frame.payload))) wfile.flush()
def handle_websockets(cls, rfile, wfile): frame = websockets.Frame.from_file(rfile) wfile.write(bytes(websockets.Frame(fin=1, opcode=frame.header.opcode, payload=frame.payload))) wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) wfile.flush() with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(rfile)
def _handle_data_frame(self, frame, source_conn, other_conn, is_server): fb = self.server_frame_buffer if is_server else self.client_frame_buffer fb.append(frame) if frame.header.fin: payload = b''.join(f.payload for f in fb) original_chunk_sizes = [len(f.payload) for f in fb] message_type = fb[0].header.opcode compressed_message = fb[0].header.rsv1 fb.clear() websocket_message = WebSocketMessage(message_type, not is_server, payload) length = len(websocket_message.content) self.flow.messages.append(websocket_message) self.channel.ask("websocket_message", self.flow) def get_chunk(payload): if len(payload) == length: # message has the same length, we can reuse the same sizes pos = 0 for s in original_chunk_sizes: yield payload[pos:pos + s] pos += s else: # just re-chunk everything into 10kB frames chunk_size = 10240 chunks = range(0, len(payload), chunk_size) for i in chunks: yield payload[i:i + chunk_size] frms = [ websockets.Frame( payload=chunk, opcode=frame.header.opcode, mask=(False if is_server else 1), masking_key=(b'' if is_server else os.urandom(4))) for chunk in get_chunk(websocket_message.content) ] if len(frms) > 0: frms[-1].header.fin = True else: frms.append( websockets.Frame( fin=True, opcode=websockets.OPCODE.CONTINUE, mask=(False if is_server else 1), masking_key=(b'' if is_server else os.urandom(4)))) frms[0].header.opcode = message_type frms[0].header.rsv1 = compressed_message for frm in frms: other_conn.send(bytes(frm)) return True
def handle_websockets(cls, rfile, wfile): wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PING, payload=b'foobar'))) wfile.flush() frame = websockets.Frame.from_file(rfile) assert frame.header.opcode == websockets.OPCODE.PONG assert frame.payload == b'foobar' wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'pong-received'))) wfile.flush()
def handle_websockets(cls, rfile, wfile): frame = websockets.Frame.from_file(rfile) assert frame.header.opcode == websockets.OPCODE.PING assert frame.payload == b'' wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PONG, payload=frame.payload))) wfile.flush() wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) wfile.flush() websockets.Frame.from_file(rfile)
def test_simple_tls(self): self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'server-foobar' self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'self.client-foobar' self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.flush()
def test_pong(self): self.setup_connection() self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.PING, payload=b'foobar'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) websockets.Frame.from_file(self.client.rfile) self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.flush() assert frame.header.opcode == websockets.OPCODE.PONG assert frame.payload == b'foobar' assert self.master.has_log("Pong Received from server", "info")
def handle_websockets(cls, rfile, wfile): wfile.write( bytes( websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'server-foobar'))) wfile.flush()
def handle_websockets(cls, rfile, wfile): frame = websockets.Frame.from_file(rfile) assert frame.header.opcode == websockets.OPCODE.TEXT success = frame.payload == b'This is an injected message!' wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=str(success).encode()))) wfile.flush()
def _handle_data_frame(self, frame, source_conn, other_conn, is_server): fb = self.server_frame_buffer if is_server else self.client_frame_buffer fb.append(frame) if frame.header.fin: if frame.header.opcode == websockets.OPCODE.TEXT: t = WebSocketTextMessage else: t = WebSocketBinaryMessage payload = b''.join(f.payload for f in fb) fb.clear() websocket_message = t(self.flow, not is_server, payload) self.flow.messages.append(websocket_message) self.channel.ask("websocket_message", self.flow) # chunk payload into multiple 10kB frames, and send them payload = websocket_message.content chunk_size = 10240 # 10kB chunks = range(0, len(payload), chunk_size) frms = [ websockets.Frame( payload=payload[i:i + chunk_size], opcode=frame.header.opcode, mask=(False if is_server else 1), masking_key=(b'' if is_server else os.urandom(4))) for i in chunks ] frms[-1].header.fin = 1 for frm in frms: other_conn.send(bytes(frm)) return True
def test_close_payload_2(self): self.setup_connection() self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar'))) self.client.wfile.flush() websockets.Frame.from_file(self.client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(self.client.rfile)
def test_close(self): client = self._setup_connection() client.wfile.write( bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) client.wfile.flush() with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile)
def test_serialization_bijection(self, masked, length): frame = websockets.Frame( os.urandom(length), fin=True, opcode=websockets.OPCODE.TEXT, mask=int(masked), masking_key=(os.urandom(4) if masked else None)) serialized = bytes(frame) assert frame == websockets.Frame.from_bytes(serialized)
def test_pong(self): client = self._setup_connection() client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PING, payload=b'foobar'))) client.wfile.flush() frame = websockets.Frame.from_file(client.rfile) assert frame.header.opcode == websockets.OPCODE.PONG assert frame.payload == b'foobar'
async def test_ping(self): self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) websockets.Frame.from_file(self.client.rfile) self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.flush() assert frame.header.opcode == websockets.OPCODE.PING assert frame.payload == b'' # We don't send payload to other end assert await self.master.await_log("Pong Received from server", "info")
def test_simple(self, streaming): class Stream: def websocket_start(self, f): f.stream = streaming self.master.addons.add(Stream()) self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'server-foobar' self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'self.client-foobar' self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef'))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) assert frame.payload == b'\xde\xad\xbe\xef' self.client.wfile.write(bytes(websockets.Frame(fin=1, mask=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.flush() assert len(self.master.state.flows) == 2 assert isinstance(self.master.state.flows[0], HTTPFlow) assert isinstance(self.master.state.flows[1], WebSocketFlow) assert len(self.master.state.flows[1].messages) == 5 assert self.master.state.flows[1].messages[0].content == 'server-foobar' assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT assert self.master.state.flows[1].messages[1].content == 'self.client-foobar' assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT assert self.master.state.flows[1].messages[2].content == 'self.client-foobar' assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef' assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef' assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
def test_ping(self): self.setup_connection() frame = websockets.Frame.from_file(self.client.rfile) assert frame.header.opcode == websockets.OPCODE.PING assert frame.payload == b'foobar' self.client.wfile.write( bytes( websockets.Frame(fin=1, opcode=websockets.OPCODE.PONG, payload=frame.payload))) self.client.wfile.flush() frame = websockets.Frame.from_file(self.client.rfile) assert frame.header.opcode == websockets.OPCODE.TEXT assert frame.payload == b'pong-received'
def test_equality(self): f = websockets.Frame(payload=b'1234') f2 = websockets.Frame(payload=b'1234') assert f == f2 assert f != b'1234'
def test_human_readable(self): f = websockets.Frame() assert repr(f) f = websockets.Frame(b"foobar") assert "foobar" in repr(f)
def handle_websockets(cls, rfile, wfile): wfile.write( bytes(websockets.Frame(fin=1, opcode=15, payload=b'foobar'))) wfile.flush()
def round(*args, **kwargs): f = websockets.Frame(*args, **kwargs) raw = bytes(f) f2 = websockets.Frame.from_file(tutils.treader(raw)) assert f == f2