Merge pull request #1127 from mitmproxy/tcp-flows

mitmdump: Add Basic Support for TCP Flows

5 files changed, 42 insertions, 12 deletions

diff --git a/test/mitmproxy/scripts/tcp_stream_modify.py b/test/mitmproxy/scripts/tcp_stream_modify.py
index 93b0d5c8..d7953ef9 100644
--- a/test/mitmproxy/scripts/tcp_stream_modify.py
+++ b/test/mitmproxy/scripts/tcp_stream_modify.py
@@ -1,3 +1,4 @@
-def tcp_message(ctx, tm):
-    if tm.sender == tm.server_conn:
-        tm.message = tm.message.replace("foo", "bar")
+def tcp_message(ctx, flow):
+    message = flow.messages[-1]
+    if not message.from_client:
+        message.content = message.content.replace("foo", "bar")
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index bf417423..3e78a5c4 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -680,6 +680,10 @@ class TestSerialize:
         for i in range(3):
             f = tutils.tflow(err=True)
             w.add(f)
+        f = tutils.ttcpflow()
+        w.add(f)
+        f = tutils.ttcpflow(err=True)
+        w.add(f)
 
         sio.seek(0)
         return flow.FlowReader(sio)
@@ -1151,6 +1155,10 @@ class TestError:
         e3 = e.copy()
         assert e3.get_state() == e.get_state()
 
+    def test_repr(self):
+        e = Error("yay")
+        assert repr(e)
+
 
 class TestClientConnection:
 
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 454736d4..0701d52b 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -14,7 +14,7 @@ from pathod import pathoc, pathod
 
 from mitmproxy.proxy.config import HostMatcher
 from mitmproxy.exceptions import Kill
-from mitmproxy.models import Error, HTTPResponse
+from mitmproxy.models import Error, HTTPResponse, HTTPFlow
 
 from . import tutils, tservers
 
@@ -177,9 +177,9 @@ class TcpMixin:
         assert n.status_code == 304
         assert i.status_code == 305
         assert i2.status_code == 306
-        assert any(f.response.status_code == 304 for f in self.master.state.flows)
-        assert not any(f.response.status_code == 305 for f in self.master.state.flows)
-        assert not any(f.response.status_code == 306 for f in self.master.state.flows)
+        assert any(f.response.status_code == 304 for f in self.master.state.flows if isinstance(f, HTTPFlow))
+        assert not any(f.response.status_code == 305 for f in self.master.state.flows if isinstance(f, HTTPFlow))
+        assert not any(f.response.status_code == 306 for f in self.master.state.flows if isinstance(f, HTTPFlow))
 
         # Test that we get the original SSL cert
         if self.ssl:
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 4fa519cc..c9d68cfd 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -50,9 +50,8 @@ class TestMaster(flow.FlowMaster):
     def clear_log(self):
         self.log = []
 
-    def handle_log(self, l):
-        self.log.append(l.msg)
-        l.reply()
+    def add_event(self, message, level=None):
+        self.log.append(message)
 
 
 class ProxyThread(threading.Thread):
diff --git a/test/mitmproxy/tutils.py b/test/mitmproxy/tutils.py
index 2dfd710e..118f849c 100644
--- a/test/mitmproxy/tutils.py
+++ b/test/mitmproxy/tutils.py
@@ -3,6 +3,8 @@ import shutil
 import tempfile
 import argparse
 import sys
+
+from mitmproxy.models.tcp import TCPMessage
 from six.moves import cStringIO as StringIO
 from contextlib import contextmanager
 
@@ -12,7 +14,7 @@ import netlib.utils
 import netlib.tutils
 from mitmproxy import utils, controller
 from mitmproxy.models import (
-    ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow
+    ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow, TCPFlow
 )
 
 
@@ -45,6 +47,26 @@ def skip_appveyor(fn):
         return fn
 
 
+def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if messages is True:
+        messages = [
+            TCPMessage(True, b"hello"),
+            TCPMessage(False, b"it's me"),
+        ]
+    if err is True:
+        err = terr()
+
+    f = TCPFlow(client_conn, server_conn)
+    f.messages = messages
+    f.error = err
+    f.reply = controller.DummyReply()
+    return f
+
+
 def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     """
     @type client_conn: bool | None | mitmproxy.proxy.connection.ClientConnection
@@ -52,7 +74,7 @@ def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     @type req:         bool | None | mitmproxy.protocol.http.HTTPRequest
     @type resp:        bool | None | mitmproxy.protocol.http.HTTPResponse
     @type err:         bool | None | mitmproxy.protocol.primitives.Error
-    @return:           bool | None | mitmproxy.protocol.http.HTTPFlow
+    @return:           mitmproxy.protocol.http.HTTPFlow
     """
     if client_conn is True:
         client_conn = tclient_conn()