From fe6f0d368b2d48d8e8e1082275e6b3e521710575 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Tue, 16 Mar 2021 15:01:33 +0100 Subject: [PATCH] inject -> inject.tcp/inject.websocket --- examples/addons/websocket-inject-message.py | 5 +- mitmproxy/addons/proxyserver.py | 56 ++++++++++++--------- test/mitmproxy/addons/test_proxyserver.py | 30 +++++++---- 3 files changed, 55 insertions(+), 36 deletions(-) diff --git a/examples/addons/websocket-inject-message.py b/examples/addons/websocket-inject-message.py index 916e2c0a9..7bc60764e 100644 --- a/examples/addons/websocket-inject-message.py +++ b/examples/addons/websocket-inject-message.py @@ -14,15 +14,16 @@ def websocket_message(flow): last_message = flow.websocket.messages[-1] if b"secret" in last_message.content: last_message.kill() - ctx.master.commands.call("inject", [flow], not last_message.from_client, "ssssssh") + ctx.master.commands.call("inject.websocket", flow, last_message.from_client, "ssssssh") # Complex example: Schedule a periodic timer async def inject_async(flow: http.HTTPFlow): msg = "hello from mitmproxy! " + assert flow.websocket # make type checker happy while flow.websocket.timestamp_end is None: - ctx.master.commands.call("inject", [flow], False, msg) + ctx.master.commands.call("inject.websocket", flow, True, msg) await asyncio.sleep(1) msg = msg[1:] + msg[:1] diff --git a/mitmproxy/addons/proxyserver.py b/mitmproxy/addons/proxyserver.py index 9ec1c176a..549d05739 100644 --- a/mitmproxy/addons/proxyserver.py +++ b/mitmproxy/addons/proxyserver.py @@ -1,9 +1,9 @@ import asyncio import warnings -from typing import Dict, Optional, Sequence, Tuple +from typing import Dict, Optional, Tuple from mitmproxy import command, controller, ctx, flow, http, log, master, options, platform, tcp, websocket -from mitmproxy.flow import Error +from mitmproxy.flow import Error, Flow from mitmproxy.proxy import commands, events from mitmproxy.proxy import server from mitmproxy.proxy.layers.tcp import TcpMessageInjected @@ -147,30 +147,36 @@ class Proxyserver: finally: del self._connections[peername] - def inject_event(self, flow: flow.Flow, event: events.Event): - if flow.client_conn.peername not in self._connections: + def inject_event(self, event: events.MessageInjected): + if event.flow.client_conn.peername not in self._connections: raise ValueError("Flow is not from a live connection.") - self._connections[flow.client_conn.peername].server_event(event) + self._connections[event.flow.client_conn.peername].server_event(event) + + @command.command("inject.websocket") + def inject_websocket(self, flow: Flow, to_client: bool, message: str, is_text: bool = True): + if not isinstance(flow, http.HTTPFlow) or not flow.websocket: + ctx.log.warn("Cannot inject WebSocket messages into non-WebSocket flows.") - @command.command("inject") - def inject(self, flows: Sequence[flow.Flow], from_client: bool, message: str): message_bytes = strutils.escaped_str_to_bytes(message) - event: events.MessageInjected - for f in flows: - if isinstance(f, http.HTTPFlow): - if f.websocket: - msg = websocket.WebSocketMessage(Opcode.TEXT, from_client, message_bytes) - event = WebSocketMessageInjected(f, msg) - else: - ctx.log.warn("Cannot inject messages into HTTP connections.") - continue - elif isinstance(f, tcp.TCPFlow): - event = TcpMessageInjected(f, tcp.TCPMessage(from_client, message_bytes)) - else: # pragma: no cover - ctx.log.warn(f"Cannot inject message into {f.__class__.__name__}, skipping.") - continue + msg = websocket.WebSocketMessage( + Opcode.TEXT if is_text else Opcode.BINARY, + not to_client, + message_bytes + ) + event = WebSocketMessageInjected(flow, msg) + try: + self.inject_event(event) + except ValueError as e: + ctx.log.warn(str(e)) - try: - self.inject_event(f, event) - except ValueError as e: - ctx.log.warn(str(e)) + @command.command("inject.tcp") + def inject_tcp(self, flow: Flow, to_client: bool, message: str): + if not isinstance(flow, tcp.TCPFlow): + ctx.log.warn("Cannot inject TCP messages into non-TCP flows.") + + message_bytes = strutils.escaped_str_to_bytes(message) + event = TcpMessageInjected(flow, tcp.TCPMessage(not to_client, message_bytes)) + try: + self.inject_event(event) + except ValueError as e: + ctx.log.warn(str(e)) diff --git a/test/mitmproxy/addons/test_proxyserver.py b/test/mitmproxy/addons/test_proxyserver.py index df62c24db..0e85054de 100644 --- a/test/mitmproxy/addons/test_proxyserver.py +++ b/test/mitmproxy/addons/test_proxyserver.py @@ -110,9 +110,9 @@ async def test_inject(): writer.write(b"a") assert await reader.read(1) == b"A" - ps.inject(state.flows, True, "b") + ps.inject_tcp(state.flows[0], False, "b") assert await reader.read(1) == b"B" - ps.inject(state.flows, False, "c") + ps.inject_tcp(state.flows[0], True, "c") assert await reader.read(1) == b"c" @@ -120,16 +120,28 @@ async def test_inject(): async def test_inject_fail(): ps = Proxyserver() with taddons.context(ps) as tctx: - ps.inject( - [tflow.tflow()], - False, + ps.inject_websocket( + tflow.tflow(), + True, "test" ) - await tctx.master.await_log("Cannot inject messages into HTTP connections.", level="warn") + await tctx.master.await_log("Cannot inject WebSocket messages into non-WebSocket flows.", level="warn") + ps.inject_tcp( + tflow.tflow(), + True, + "test" + ) + await tctx.master.await_log("Cannot inject TCP messages into non-TCP flows.", level="warn") - ps.inject( - [tflow.twebsocketflow()], - False, + ps.inject_websocket( + tflow.twebsocketflow(), + True, + "test" + ) + await tctx.master.await_log("Flow is not from a live connection.", level="warn") + ps.inject_websocket( + tflow.ttcpflow(), + True, "test" ) await tctx.master.await_log("Flow is not from a live connection.", level="warn")