Add Inline Script Hooks to TCP mode

docs/scripting/inlinescripts.rst

@@ -36,14 +36,13 @@ We encourage you to either browse them locally or on `GitHub`_.
 Events
 ------
 
-.. TODO: Split this into Connection, HTTP and TCP events once we have TCP events.
-
 The ``context`` argument passed to each event method is always a
 :py:class:`~libmproxy.script.ScriptContext` instance. It is guaranteed to be the same object
 for the scripts lifetime and is not shared between multiple inline scripts. You can safely use it
 to store any form of state you require.
 
-Events are listed in the order they usually occur.
+Script Lifecycle Events
+^^^^^^^^^^^^^^^^^^^^^^^
 
 .. py:function:: start(context, argv)
 
@@ -52,6 +51,13 @@ Events are listed in the order they usually occur.
     :param List[str] argv: The inline scripts' arguments.
         For example, ``mitmproxy -s 'example.py --foo 42'`` sets argv to ``["--foo", "42"]``.
 
+.. py:function:: done(context)
+
+    Called once on script shutdown, after any other events.
+
+Connection Events
+^^^^^^^^^^^^^^^^^
+
 .. py:function:: clientconnect(context, root_layer)
 
     Called when a client initiates a connection to the proxy. Note that
@@ -64,14 +70,13 @@ Events are listed in the order they usually occur.
         :py:class:`~libmproxy.proxy.RootContext`. For example, ``root_layer.client_conn.address``
         gives the remote address of the connecting client.
 
+.. py:function:: clientdisconnect(context, root_layer)
 
-.. py:function:: request(context, flow)
+    Called when a client disconnects from the proxy.
 
-    Called when a client request has been received. The ``flow`` object is
+    .. versionchanged:: 0.14
-    guaranteed to have a non-None ``request`` attribute.
 
-    :param HTTPFlow flow: The flow containing the request which has been received.
+    :param Layer root_layer: see :py:func:`clientconnect`
-        The object is guaranteed to have a non-None ``request`` attribute.
 
 .. py:function:: serverconnect(context, server_conn)
 
@@ -81,6 +86,25 @@ Events are listed in the order they usually occur.
     :param ServerConnection server_conn: The server connection object. It is guaranteed to have a
         non-None ``address`` attribute.
 
+.. py:function:: serverdisconnect(context, server_conn)
+
+    Called when the proxy has closed the server connection.
+
+    .. versionadded:: 0.14
+
+    :param ServerConnection server_conn: see :py:func:`serverconnect`
+
+HTTP Events
+^^^^^^^^^^^
+
+.. py:function:: request(context, flow)
+
+    Called when a client request has been received. The ``flow`` object is
+    guaranteed to have a non-None ``request`` attribute.
+
+    :param HTTPFlow flow: The flow containing the request which has been received.
+        The object is guaranteed to have a non-None ``request`` attribute.
+
 .. py:function:: responseheaders(context, flow)
 
     Called when the headers of a server response have been received.
@@ -109,26 +133,19 @@ Events are listed in the order they usually occur.
     :param HTTPFlow flow: The flow containing the error.
         It is guaranteed to have non-None ``error`` attribute.
 
-.. py:function:: serverdisconnect(context, server_conn)
+TCP Events
+^^^^^^^^^^
 
-    Called when the proxy has closed the server connection.
+.. py:function:: tcp_message(context, tcp_msg)
 
-    .. versionadded:: 0.14
+    .. warning::  API is subject to change
 
-    :param ServerConnection server_conn: see :py:func:`serverconnect`
+    If the proxy is in :ref:`TCP mode <tcpproxy>`, this event is called when it
+    receives a TCP payload from the client or server.
 
-.. py:function:: clientdisconnect(context, root_layer)
+    The sender and receiver are identifiable. The message is user-modifiable.
-
-    Called when a client disconnects from the proxy.
-
-    .. versionchanged:: 0.14
-
-    :param Layer root_layer: see :py:func:`clientconnect`
-
-.. py:function:: done(context)
-
-    Called once on script shutdown, after any other events.
 
+    :param TcpMessage tcp_msg: see *examples/tcp_message.py*
 
 API
 ---

examples/tcp_message.py

@@ -0,0 +1,24 @@
+'''
+tcp_message Inline Script Hook API Demonstration
+------------------------------------------------
+
+* modifies packets containing "foo" to "bar"
+* prints various details for each packet.
+
+example cmdline invocation:
+mitmdump -T --host --tcp ".*" -q -s examples/tcp_message.py
+'''
+from netlib.utils import clean_bin
+
+def tcp_message(ctx, tcp_msg):
+    modified_msg = tcp_msg.message.replace("foo", "bar")
+
+    is_modified = False if modified_msg == tcp_msg.message else True
+    tcp_msg.message = modified_msg
+
+    print("[tcp_message{}] from {} {} to {} {}:\r\n{}".format(
+        " (modified)" if is_modified else "",
+        "client" if tcp_msg.sender == tcp_msg.client_conn else "server",
+        tcp_msg.sender.address,
+        "server" if tcp_msg.receiver == tcp_msg.server_conn else "client",
+        tcp_msg.receiver.address, clean_bin(tcp_msg.message)))

libmproxy/flow.py

@@ -1050,6 +1050,10 @@ class FlowMaster(controller.Master):
             self.add_event('"{}" reloaded.'.format(s.filename))
         return ok
 
+    def handle_tcp_message(self, m):
+        self.run_script_hook("tcp_message", m)
+        m.reply()
+
     def shutdown(self):
         self.unload_scripts()
         controller.Master.shutdown(self)

libmproxy/protocol/rawtcp.py

@@ -13,6 +13,15 @@ from ..exceptions import ProtocolException
 from .base import Layer
 
 
+class TcpMessage(object):
+    def __init__(self, client_conn, server_conn, sender, receiver, message):
+        self.client_conn = client_conn
+        self.server_conn = server_conn
+        self.sender = sender
+        self.receiver = receiver
+        self.message = message
+
+
 class RawTCPLayer(Layer):
     chunk_size = 4096
 
@@ -50,7 +59,13 @@ class RawTCPLayer(Layer):
                             return
                         continue
 
-                    dst.sendall(buf[:size])
+                    tcp_message = TcpMessage(
+                        self.client_conn, self.server_conn,
+                        self.client_conn if dst == server else self.server_conn,
+                        self.server_conn if dst == server else self.client_conn,
+                        buf[:size].tobytes())
+                    self.channel.ask("tcp_message", tcp_message)
+                    dst.sendall(tcp_message.message)
 
                     if self.logging:
                         # log messages are prepended with the client address,
@@ -59,7 +74,7 @@ class RawTCPLayer(Layer):
                             direction = "-> tcp -> {}".format(repr(self.server_conn.address))
                         else:
                             direction = "<- tcp <- {}".format(repr(self.server_conn.address))
-                        data = clean_bin(buf[:size].tobytes())
+                        data = clean_bin(tcp_message.message)
                         self.log(
                             "{}\r\n{}".format(direction, data),
                             "info"

test/scripts/tcp_stream_modify.py

@@ -0,0 +1,3 @@
+def tcp_message(ctx,tm):
+    if tm.sender == tm.server_conn:
+        tm.message = tm.message.replace("foo", "bar")

test/test_server.py

@@ -502,6 +502,18 @@ class TestHttps2Http(tservers.ReverseProxTest):
 class TestTransparent(tservers.TransparentProxTest, CommonMixin, TcpMixin):
     ssl = False
 
+    def test_tcp_stream_modify(self):
+        self.master.load_script(
+            tutils.test_data.path("scripts/tcp_stream_modify.py"))
+
+        self._tcpproxy_on()
+        d = self.pathod('200:b"foo"')
+        self._tcpproxy_off()
+
+        assert d.content == "bar"
+
+        self.master.unload_scripts()
+
 
 class TestTransparentSSL(tservers.TransparentProxTest, CommonMixin, TcpMixin):
     ssl = True