add TCP message injection

mitmproxy/addons/proxyserver.py

@@ -6,8 +6,9 @@ from mitmproxy import command, controller, ctx, flow, http, log, master, options
 from mitmproxy.flow import Error
 from mitmproxy.proxy import commands, events
 from mitmproxy.proxy import server
+from mitmproxy.proxy.layers.tcp import TcpMessageInjected
 from mitmproxy.proxy.layers.websocket import WebSocketMessageInjected
-from mitmproxy.utils import asyncio_utils, human
+from mitmproxy.utils import asyncio_utils, human, strutils
 from wsproto.frame_protocol import Opcode
 
 
@@ -151,24 +152,23 @@ class Proxyserver:
             raise ValueError("Flow is not from a live connection.")
         self._connections[flow.client_conn.peername].server_event(event)
 
-    @command.command("inject.text")
+    @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:
-                    event = WebSocketMessageInjected(
+                    msg = websocket.WebSocketMessage(Opcode.TEXT, from_client, message_bytes)
-                        f,
+                    event = WebSocketMessageInjected(f, msg)
-                        websocket.WebSocketMessage(
-                            Opcode.TEXT, from_client, message.encode()
-                        )
-                    )
                 else:
                     ctx.log.warn("Cannot inject messages into HTTP connections.")
                     continue
             elif isinstance(f, tcp.TCPFlow):
-                raise NotImplementedError
+                event = TcpMessageInjected(f, tcp.TCPMessage(from_client, message_bytes))
             else:
-                raise NotImplementedError
+                ctx.log.warn(f"Cannot inject message into {f.__class__.__name__}, skipping.")
+
             try:
                 self.inject_event(f, event)
             except ValueError as e:

mitmproxy/proxy/layers/http/__init__.py

@@ -589,14 +589,22 @@ class HttpLayer(layer.Layer):
             yield from self.event_to_child(stream, event)
         elif isinstance(event, events.MessageInjected):
             # For injected messages we pass the HTTP stacks entirely and directly address the stream.
-            conn = self.connections[event.flow.server_conn]
+            try:
+                conn = self.connections[event.flow.server_conn]
+            except KeyError:
+                # We have a miss for the server connection, which means we're looking at a connection object
+                # that is tunneled over another connection (for example: over an upstream HTTP proxy).
+                # We now take the stream associated with the client connection. That won't work for HTTP/2,
+                # but it's good enough for HTTP/1.
+                conn = self.connections[event.flow.client_conn]
             if isinstance(conn, HttpStream):
                 stream_id = conn.stream_id
             else:
                 # We reach to the end of the connection's child stack to get the HTTP/1 client layer,
                 # which tells us which stream we are dealing with.
                 conn = conn.context.layers[-1]
-                assert isinstance(conn, Http1Client)
+                assert isinstance(conn, Http1Connection)
+                assert conn.stream_id
                 stream_id = conn.stream_id
             yield from self.event_to_child(self.streams[stream_id], event)
         elif isinstance(event, events.ConnectionEvent):

mitmproxy/proxy/layers/tcp.py

@@ -6,6 +6,7 @@ from mitmproxy.proxy import commands, events, layer
 from mitmproxy.proxy.commands import StartHook
 from mitmproxy.connection import ConnectionState, Connection
 from mitmproxy.proxy.context import Context
+from mitmproxy.proxy.events import MessageInjected
 from mitmproxy.proxy.utils import expect
 
 
@@ -45,6 +46,12 @@ class TcpErrorHook(StartHook):
     flow: tcp.TCPFlow
 
 
+class TcpMessageInjected(MessageInjected[tcp.TCPMessage]):
+    """
+    The user has injected a custom TCP message.
+    """
+
+
 class TCPLayer(layer.Layer):
     """
     Simple TCP layer that just relays messages right now.
@@ -76,8 +83,18 @@ class TCPLayer(layer.Layer):
 
     _handle_event = start
 
-    @expect(events.DataReceived, events.ConnectionClosed)
+    @expect(events.DataReceived, events.ConnectionClosed, TcpMessageInjected)
-    def relay_messages(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
+    def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:
+
+        if isinstance(event, TcpMessageInjected):
+            # we just spoof that we received data here and then process that regularly.
+            event = events.DataReceived(
+                self.context.client if event.message.from_client else self.context.server,
+                event.message.content,
+            )
+
+        assert isinstance(event, events.ConnectionEvent)
+
         from_client = event.connection == self.context.client
         send_to: Connection
         if from_client:
@@ -110,6 +127,8 @@ class TCPLayer(layer.Layer):
                     yield TcpEndHook(self.flow)
             else:
                 yield commands.CloseConnection(send_to, half_close=True)
+        else:
+            raise AssertionError(f"Unexpected event: {event}")
 
     @expect(events.DataReceived, events.ConnectionClosed)
     def done(self, _) -> layer.CommandGenerator[None]:

test/mitmproxy/proxy/layers/http/test_http.py

@@ -9,9 +9,9 @@ from mitmproxy.proxy.commands import CloseConnection, OpenConnection, SendData,
 from mitmproxy.connection import ConnectionState, Server
 from mitmproxy.proxy.events import ConnectionClosed, DataReceived
 from mitmproxy.proxy.layers import TCPLayer, http, tls
-from mitmproxy.proxy.layers.tcp import TcpStartHook
+from mitmproxy.proxy.layers.tcp import TcpMessageInjected, TcpStartHook
 from mitmproxy.proxy.layers.websocket import WebsocketStartHook
-from mitmproxy.tcp import TCPFlow
+from mitmproxy.tcp import TCPFlow, TCPMessage
 from test.mitmproxy.proxy.tutils import Placeholder, Playbook, reply, reply_next_layer
 
 
@@ -543,6 +543,7 @@ def test_upstream_proxy(tctx, redirect, scheme):
 def test_http_proxy_tcp(tctx, mode, close_first):
     """Test TCP over HTTP CONNECT."""
     server = Placeholder(Server)
+    f = Placeholder(TCPFlow)
 
     if mode == "upstream":
         tctx.options.mode = "upstream:http://proxy:8080"
@@ -558,7 +559,9 @@ def test_http_proxy_tcp(tctx, mode, close_first):
             << SendData(tctx.client, b"HTTP/1.1 200 Connection established\r\n\r\n")
             >> DataReceived(tctx.client, b"this is not http")
             << layer.NextLayerHook(Placeholder())
-            >> reply_next_layer(lambda ctx: TCPLayer(ctx, ignore=True))
+            >> reply_next_layer(lambda ctx: TCPLayer(ctx, ignore=False))
+            << TcpStartHook(f)
+            >> reply()
             << OpenConnection(server)
     )
 
@@ -579,6 +582,12 @@ def test_http_proxy_tcp(tctx, mode, close_first):
     else:
         assert server().address == ("proxy", 8080)
 
+    assert (
+        playbook
+        >> TcpMessageInjected(f, TCPMessage(False, b"fake news from your friendly man-in-the-middle"))
+        << SendData(tctx.client, b"fake news from your friendly man-in-the-middle")
+    )
+
     if close_first == "client":
         a, b = tctx.client, server
     else:

test/mitmproxy/proxy/layers/test_tcp.py

@@ -4,7 +4,8 @@ from mitmproxy.proxy.commands import CloseConnection, OpenConnection, SendData
 from mitmproxy.connection import ConnectionState
 from mitmproxy.proxy.events import ConnectionClosed, DataReceived
 from mitmproxy.proxy.layers import tcp
-from mitmproxy.tcp import TCPFlow
+from mitmproxy.proxy.layers.tcp import TcpMessageInjected
+from mitmproxy.tcp import TCPFlow, TCPMessage
 from ..tutils import Placeholder, Playbook, reply
 
 
@@ -122,3 +123,27 @@ def test_ignore(tctx, ignore):
     else:
         with pytest.raises(AssertionError):
             no_flow_hooks()
+
+
+def test_inject(tctx):
+    """inject data into an open connection."""
+    f = Placeholder(TCPFlow)
+
+    assert (
+            Playbook(tcp.TCPLayer(tctx))
+            << tcp.TcpStartHook(f)
+            >> TcpMessageInjected(f, TCPMessage(True, b"hello!"))
+            >> reply(to=-2)
+            << OpenConnection(tctx.server)
+            >> reply(None)
+            << tcp.TcpMessageHook(f)
+            >> reply()
+            << SendData(tctx.server, b"hello!")
+            # and the other way...
+            >> TcpMessageInjected(f, TCPMessage(False, b"I have already done the greeting for you."))
+            << tcp.TcpMessageHook(f)
+            >> reply()
+            << SendData(tctx.client, b"I have already done the greeting for you.")
+            << None
+    )
+    assert len(f().messages) == 2