inject -> inject.tcp/inject.websocket

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]
 

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))

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")