diff options
Diffstat (limited to 'test')
-rw-r--r-- | test/mitmproxy/proxy/protocol/test_websocket.py | 43 | ||||
-rw-r--r-- | test/mitmproxy/test_websocket.py | 12 |
2 files changed, 55 insertions, 0 deletions
diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py index 1f4e2bca..0b26ed29 100644 --- a/test/mitmproxy/proxy/protocol/test_websocket.py +++ b/test/mitmproxy/proxy/protocol/test_websocket.py @@ -467,3 +467,46 @@ class TestExtension(_WebSocketTest): 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 + + +class TestInjectMessageClient(_WebSocketTest): + + @classmethod + def handle_websockets(cls, rfile, wfile): + pass + + def test_inject_message_client(self): + class Inject: + def websocket_start(self, flow): + flow.inject_message(flow.client_conn, 'This is an injected message!') + + self.proxy.set_addons(Inject()) + self.setup_connection() + + frame = websockets.Frame.from_file(self.client.rfile) + assert frame.header.opcode == websockets.OPCODE.TEXT + assert frame.payload == b'This is an injected message!' + + +class TestInjectMessageServer(_WebSocketTest): + + @classmethod + 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 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() + + frame = websockets.Frame.from_file(self.client.rfile) + assert frame.header.opcode == websockets.OPCODE.TEXT + assert frame.payload == b'True' diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py index fcacec36..bd4bb4c9 100644 --- a/test/mitmproxy/test_websocket.py +++ b/test/mitmproxy/test_websocket.py @@ -92,3 +92,15 @@ class TestWebSocketFlow: assert not f.messages[-1].killed f.messages[-1].kill() assert f.messages[-1].killed + + def test_inject_message(self): + f = tflow.twebsocketflow() + + with pytest.raises(ValueError): + f.inject_message(None, 'foobar') + + f.inject_message(f.client_conn, 'foobar') + assert f._inject_messages_client.qsize() == 1 + + f.inject_message(f.server_conn, 'foobar') + assert f._inject_messages_client.qsize() == 1 |