def test_extension(self): self.setup_connection(True) header, _, _ = websocket_utils.read_raw_frame(self.client.rfile) assert header.rsv.rsv1 self.client.wfile.write(b'\xc1\x8fQ\xb7vX\x1by\xbf\x14\x9c\x9c\xa7\x15\x9ax9\x12}\xb5v') self.client.wfile.flush() header, _, _ = websocket_utils.read_raw_frame(self.client.rfile) assert header.rsv.rsv1 self.client.wfile.write(b'\xc2\x87\xeb\xbb\x0csQ\x0cz\xac\x90\xbb\x0c') self.client.wfile.flush() header, _, _ = websocket_utils.read_raw_frame(self.client.rfile) assert header.rsv.rsv1 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 == Opcode.TEXT assert self.master.state.flows[1].messages[1].content == 'client-foobar' assert self.master.state.flows[1].messages[1].type == Opcode.TEXT assert self.master.state.flows[1].messages[2].content == 'client-foobar' assert self.master.state.flows[1].messages[2].type == 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 == 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 == Opcode.BINARY
def handle_websockets(cls, rfile, wfile): header, frame, _ = websocket_utils.read_raw_frame(rfile) wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=header.opcode, payload=frame.payload))) wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.CLOSE))) wfile.flush() with pytest.raises(exceptions.TcpDisconnect): _, _, _ = websocket_utils.read_raw_frame(rfile)
def test_close_payload_2(self): self.setup_connection() self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE, payload=b'\00\42foobar'))) self.client.wfile.flush() _ = websocket_utils.read_raw_frame(self.client.rfile) with pytest.raises(exceptions.TcpDisconnect): _ = websocket_utils.read_raw_frame(self.client.rfile)
def handle_websockets(cls, rfile, wfile): wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.TEXT, payload=b'server-foobar'))) wfile.flush() header, frame, _ = websocket_utils.read_raw_frame(rfile) wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=header.opcode, payload=frame.payload))) wfile.flush() header, frame, _ = websocket_utils.read_raw_frame(rfile) wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=header.opcode, payload=frame.payload))) wfile.flush()
def handle_websockets(cls, rfile, wfile): header, frame, _ = websocket_utils.read_raw_frame(rfile) assert header.opcode == Opcode.PING assert frame.payload == b'' wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.PONG, payload=frame.payload))) wfile.flush() wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.CLOSE))) wfile.flush() _ = websocket_utils.read_raw_frame(rfile)
async def test_ping(self): self.setup_connection() header, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) _ = websocket_utils.read_raw_frame(self.client.rfile) self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE))) self.client.wfile.flush() assert header.opcode == 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 handle_websockets(cls, rfile, wfile): wfile.write(b'\xc1\x0f*N-*K-\xd2M\xcb\xcfOJ,\x02\x00') wfile.flush() header, _, _ = websocket_utils.read_raw_frame(rfile) assert header.rsv.rsv1 wfile.write(b'\xc1\nJ\xce\xc9L\xcd+\x81r\x00\x00') wfile.flush() header, _, _ = websocket_utils.read_raw_frame(rfile) assert header.rsv.rsv1 wfile.write(b'\xc2\x07\xba\xb7v\xdf{\x00\x00') wfile.flush()
async def test_pong(self): self.setup_connection() self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.PING, payload=b'foobar'))) self.client.wfile.flush() header, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) _ = websocket_utils.read_raw_frame(self.client.rfile) self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE))) self.client.wfile.flush() assert header.opcode == Opcode.PONG assert frame.payload == b'foobar' assert await self.master.await_log("pong received")
def test_simple_tls(self): self.setup_connection() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'server-foobar' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.TEXT, payload=b'self.client-foobar'))) self.client.wfile.flush() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'self.client-foobar' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.CLOSE))) self.client.wfile.flush()
def handle_websockets(cls, rfile, wfile): header, frame, _ = websocket_utils.read_raw_frame(rfile) assert header.opcode == Opcode.TEXT success = frame.payload == b'This is an injected message!' wfile.write(bytes(websockets_frame.Frame(fin=1, opcode=Opcode.TEXT, payload=str(success).encode()))) wfile.flush()
def test_invalid_frame(self): self.setup_connection() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) code, = struct.unpack('!H', frame.payload[:2]) assert code == 1002 assert frame.payload[2:].startswith(b'Invalid opcode')
def test_streaming(self, streaming): class Stream: def websocket_start(self, f): f.stream = streaming self.proxy.set_addons(Stream()) self.setup_connection() frame = None if not streaming: with pytest.raises(exceptions.TcpDisconnect): # Reader.safe_read get nothing as result _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame is None else: _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame assert self.master.state.flows[1].messages == [] # Message not appended as the final frame isn't received
def test_kill(self): class KillFlow: def websocket_message(self, f): f.kill() self.proxy.set_addons(KillFlow()) self.setup_connection() with pytest.raises(exceptions.TcpDisconnect): _, _, _ = websocket_utils.read_raw_frame(self.client.rfile)
def test_inject_message_server(self): class Inject: def websocket_start(self, flow): flow.inject_message(flow.server_conn, 'This is an injected message!') self.proxy.set_addons(Inject()) self.setup_connection() header, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert header.opcode == Opcode.TEXT assert frame.payload == b'True'
def __call__(self): self.flow = WebSocketFlow(self.client_conn, self.server_conn, self.handshake_flow) self.flow.metadata['websocket_handshake'] = self.handshake_flow.id self.handshake_flow.metadata['websocket_flow'] = self.flow.id self.channel.ask("websocket_start", self.flow) conns = [c.connection for c in self.connections.keys()] close_received = False try: while not self.channel.should_exit.is_set(): self._inject_messages(self.client_conn, self.flow._inject_messages_client) self._inject_messages(self.server_conn, self.flow._inject_messages_server) r = tcp.ssl_read_select(conns, 0.1) for conn in r: source_conn = self.client_conn if conn == self.client_conn.connection else self.server_conn other_conn = self.server_conn if conn == self.client_conn.connection else self.client_conn is_server = (source_conn == self.server_conn) header, frame, consumed_bytes = websocket_utils.read_raw_frame( source_conn.rfile) self.log( "WebSocket Frame from {}: {}, {}".format( "server" if is_server else "client", header, frame, ), "debug") data = self.connections[source_conn].receive_data( consumed_bytes) source_conn.send(data) if close_received: return for event in self.connections[source_conn].events(): if not self._handle_event(event, source_conn, other_conn, is_server): if not close_received: close_received = True except (socket.error, exceptions.TcpException, SSL.Error) as e: s = 'server' if is_server else 'client' self.flow.error = flow.Error( "WebSocket connection closed unexpectedly by {}: {}".format( s, repr(e))) self.channel.tell("websocket_error", self.flow) finally: self.flow.ended = True self.channel.tell("websocket_end", self.flow)
def test_simple(self, streaming): class Stream: def websocket_start(self, f): f.stream = streaming self.proxy.set_addons(Stream()) self.setup_connection() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'server-foobar' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.TEXT, payload=b'self.client-foobar'))) self.client.wfile.flush() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'self.client-foobar' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.BINARY, payload=b'\xde\xad\xbe\xef'))) self.client.wfile.flush() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'\xde\xad\xbe\xef' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=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 == Opcode.TEXT assert self.master.state.flows[1].messages[1].content == 'self.client-foobar' assert self.master.state.flows[1].messages[1].type == Opcode.TEXT assert self.master.state.flows[1].messages[2].content == 'self.client-foobar' assert self.master.state.flows[1].messages[2].type == 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 == 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 == 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, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'foo' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.TEXT, payload=b'self.client-foobar'))) self.client.wfile.flush() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'foo' self.client.wfile.write(bytes(websockets_frame.Frame(fin=1, mask=1, opcode=Opcode.BINARY, payload=b'\xde\xad\xbe\xef'))) self.client.wfile.flush() _, frame, _ = websocket_utils.read_raw_frame(self.client.rfile) assert frame.payload == b'foo'
def test_read_raw_frame(input, masking_key, payload_length): bio = BytesIO(input) bio.safe_read = bio.read header, frame, consumed_bytes = websocket_utils.read_raw_frame(bio) assert header == \ Header( fin=False, rsv=RsvBits(rsv1=False, rsv2=False, rsv3=False), opcode=Opcode.TEXT, payload_len=payload_length, masking_key=masking_key, ) assert frame == \ Frame( opcode=Opcode.TEXT, payload=b'server-foobar', frame_finished=False, message_finished=False, ) assert consumed_bytes == input