From fa7246279817b7c01448fd4059c8d2be34e84f8b Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Fri, 1 Jul 2016 17:17:16 -0700 Subject: fix tcp message handling --- mitmproxy/console/master.py | 14 +++++++++++++- mitmproxy/contentviews.py | 2 +- mitmproxy/dump.py | 21 ++++++++++++++++++--- mitmproxy/flow/master.py | 9 +-------- netlib/http/request.py | 4 +++- netlib/strutils.py | 7 ++++--- test/mitmproxy/test_server.py | 7 +++---- test/mitmproxy/tservers.py | 2 +- 8 files changed, 44 insertions(+), 22 deletions(-) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 5fd51f4b..95c9704d 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -30,7 +30,7 @@ from mitmproxy.console import palettes from mitmproxy.console import signals from mitmproxy.console import statusbar from mitmproxy.console import window -from netlib import tcp +from netlib import tcp, strutils EVENTLOG_SIZE = 500 @@ -797,6 +797,18 @@ class ConsoleMaster(flow.FlowMaster): self.process_flow(f) return f + @controller.handler + def tcp_message(self, f): + super(ConsoleMaster, self).tcp_message(f) + message = f.messages[-1] + direction = "->" if message.from_client else "<-" + self.add_event("{client} {direction} tcp {direction} {server}".format( + client=repr(f.client_conn.address), + server=repr(f.server_conn.address), + direction=direction, + ), "info") + self.add_event(strutils.bytes_to_escaped_str(message.content), "debug") + @controller.handler def script_change(self, script): if super(ConsoleMaster, self).script_change(script): diff --git a/mitmproxy/contentviews.py b/mitmproxy/contentviews.py index 7c9e4ba1..de88c9ea 100644 --- a/mitmproxy/contentviews.py +++ b/mitmproxy/contentviews.py @@ -600,7 +600,7 @@ def safe_to_print(lines, encoding="utf8"): try: text = strutils.clean_bin(text.decode(encoding, "strict")) except UnicodeDecodeError: - text = strutils.clean_bin(text).decode(encoding, "strict") + text = strutils.clean_bin(text) clean_line.append((style, text)) yield clean_line diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index cc6896ed..6670be9b 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -178,7 +178,7 @@ class DumpMaster(flow.FlowMaster): click.secho(text, file=self.outfile, **style) def _echo_message(self, message): - if self.o.flow_detail >= 2: + if self.o.flow_detail >= 2 and hasattr(message, "headers"): headers = "\r\n".join( "{}: {}".format( click.style(strutils.bytes_to_escaped_str(k), fg="blue", bold=True), @@ -196,7 +196,7 @@ class DumpMaster(flow.FlowMaster): type, lines = contentviews.get_content_view( contentviews.get("Auto"), message.content, - headers=message.headers + headers=getattr(message, "headers", None) ) except exceptions.ContentViewException: s = "Content viewer failed: \n" + traceback.format_exc() @@ -204,7 +204,7 @@ class DumpMaster(flow.FlowMaster): type, lines = contentviews.get_content_view( contentviews.get("Raw"), message.content, - headers=message.headers + headers=getattr(message, "headers", None) ) styles = dict( @@ -352,6 +352,21 @@ class DumpMaster(flow.FlowMaster): self._process_flow(f) return f + @controller.handler + def tcp_message(self, f): + super(DumpMaster, self).tcp_message(f) + + if self.o.flow_detail == 0: + return + message = f.messages[-1] + direction = "->" if message.from_client else "<-" + self.echo("{client} {direction} tcp {direction} {server}".format( + client=repr(f.client_conn.address), + server=repr(f.server_conn.address), + direction=direction, + )) + self._echo_message(message) + def run(self): # pragma: no cover if self.o.rfile and not self.o.keepserving: self.unload_scripts() # make sure to trigger script unload events. diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 289102a1..efb5d013 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -499,15 +499,8 @@ class FlowMaster(controller.Master): @controller.handler def tcp_message(self, flow): + # type: (TCPFlow) -> None self.run_scripts("tcp_message", flow) - message = flow.messages[-1] - direction = "->" if message.from_client else "<-" - self.add_event("{client} {direction} tcp {direction} {server}".format( - client=repr(flow.client_conn.address), - server=repr(flow.server_conn.address), - direction=direction, - ), "info") - self.add_event(strutils.clean_bin(message.content), "debug") @controller.handler def tcp_error(self, flow): diff --git a/netlib/http/request.py b/netlib/http/request.py index ff057b79..d9f4ed00 100644 --- a/netlib/http/request.py +++ b/netlib/http/request.py @@ -28,7 +28,7 @@ class RequestData(message.MessageData): self.first_line_format = first_line_format self.method = method - self.scheme = scheme or b'' + self.scheme = scheme self.host = host self.port = port self.path = path @@ -106,6 +106,8 @@ class Request(message.Message): """ HTTP request scheme, which should be "http" or "https". """ + if not self.data.scheme: + return self.data.scheme return message._native(self.data.scheme) @scheme.setter diff --git a/netlib/strutils.py b/netlib/strutils.py index ca6eaa42..414b2e57 100644 --- a/netlib/strutils.py +++ b/netlib/strutils.py @@ -29,6 +29,7 @@ def native(s, *encoding_opts): def clean_bin(s, keep_spacing=True): + # type: (Union[bytes, six.text_type], bool) -> six.text_type """ Cleans binary data to make it safe to display. @@ -49,8 +50,8 @@ def clean_bin(s, keep_spacing=True): keep = (9, 10, 13) # \t, \n, \r, else: keep = () - return b"".join( - six.int2byte(ch) if (31 < ch < 127 or ch in keep) else b"." + return "".join( + chr(ch) if (31 < ch < 127 or ch in keep) else "." for ch in six.iterbytes(s) ) @@ -133,6 +134,6 @@ def hexdump(s): for i in range(0, len(s), 16): offset = "{:0=10x}".format(i).encode() part = s[i:i + 16] - x = b" ".join("{:0=2x}".format(i).encode() for i in six.iterbytes(part)) + x = " ".join("{:0=2x}".format(i) for i in six.iterbytes(part)) x = x.ljust(47) # 16*2 + 15 yield (offset, x, clean_bin(part, False)) diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index c5caa21f..1bbef975 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -190,10 +190,9 @@ class TcpMixin: assert i_cert == i2_cert == n_cert # Make sure that TCP messages are in the event log. - # print(m for m in self.master.tlog) - # print(self.master.tlog) - assert any("305" in m for m in self.master.tlog) - assert any("306" in m for m in self.master.tlog) + # Re-enable and fix this when we start keeping TCPFlows in the state. + # assert any("305" in m for m in self.master.tlog) + # assert any("306" in m for m in self.master.tlog) class AppMixin: diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index d9855f9a..51f4b4e2 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -45,7 +45,7 @@ class TestMaster(flow.FlowMaster): self.tlog = [] def add_event(self, message, level=None): - self.tlog.append(strutils.native(message, "utf8")) + self.tlog.append(message) class ProxyThread(threading.Thread): -- cgit v1.2.3