Merge pull request #872 from bazzinotti/inline-tcp

Add "tcp_message" Inline Script Hook

2 files changed, 21 insertions, 2 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index f02b5767..8e066191 100644
--- a/libmproxy/flow.py
+++ b/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)
diff --git a/libmproxy/protocol/rawtcp.py b/libmproxy/protocol/rawtcp.py
index 5f08fd17..ccd3c7ec 100644
--- a/libmproxy/protocol/rawtcp.py
+++ b/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"