minor fixes, adjust tests

mitmproxy/models/tcp.py

@@ -34,7 +34,7 @@ class TCPMessage(Serializable):
 
 class TCPFlow(Flow):
     """
-    A SSHFlow is a simplified representation of an SSH session.
+    A TCPFlow is a simplified representation of a TCP session.
     """
 
     def __init__(self, client_conn, server_conn, live=None):

mitmproxy/protocol/rawtcp.py

@@ -48,12 +48,12 @@ class RawTCPLayer(Layer):
                         if isinstance(conn, SSL.Connection):
                             # We can't half-close a connection, so we just close everything here.
                             # Sockets will be cleaned up on a higher level.
-                            break
+                            return
                         else:
                             dst.shutdown(socket.SHUT_WR)
 
                         if len(conns) == 0:
-                            break
+                            return
                         continue
 
                     tcp_message = TCPMessage(dst == server, buf[:size].tobytes())

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

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:

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