From acd51befbb89f28397f8e3e52ead5bfa11fdc93f Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 11 May 2016 11:13:57 -0600 Subject: [PATCH] minor fixes, adjust tests --- mitmproxy/models/tcp.py | 2 +- mitmproxy/protocol/rawtcp.py | 4 ++-- test/mitmproxy/scripts/tcp_stream_modify.py | 7 ++++--- test/mitmproxy/test_server.py | 8 ++++---- test/mitmproxy/tservers.py | 5 ++--- 5 files changed, 13 insertions(+), 13 deletions(-) diff --git a/mitmproxy/models/tcp.py b/mitmproxy/models/tcp.py index 22fd1f0c8..7e966b95c 100644 --- a/mitmproxy/models/tcp.py +++ b/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): diff --git a/mitmproxy/protocol/rawtcp.py b/mitmproxy/protocol/rawtcp.py index 5f6fca753..1b546c40d 100644 --- a/mitmproxy/protocol/rawtcp.py +++ b/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()) diff --git a/test/mitmproxy/scripts/tcp_stream_modify.py b/test/mitmproxy/scripts/tcp_stream_modify.py index 93b0d5c8d..d7953ef99 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_server.py b/test/mitmproxy/test_server.py index 454736d4d..0701d52bb 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 4fa519cc6..c9d68cfd0 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):