From 43c52054241d8743996e8fd5321db93ada37a7b1 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Tue, 10 May 2016 14:03:14 -0600 Subject: [PATCH 1/6] mimtdump: add basic support for tcp flows --- mitmproxy/flow.py | 61 ++++++++++++++++++++++++++++----- mitmproxy/models/__init__.py | 5 ++- mitmproxy/models/flow.py | 9 +++++ mitmproxy/models/http.py | 6 ---- mitmproxy/models/tcp.py | 50 +++++++++++++++++++++++++++ mitmproxy/protocol/rawtcp.py | 61 ++++++++++++--------------------- mitmproxy/proxy/root_context.py | 2 +- 7 files changed, 139 insertions(+), 55 deletions(-) create mode 100644 mitmproxy/models/tcp.py diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index ccedd1d42..1d05d4bb9 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -3,7 +3,6 @@ """ from __future__ import absolute_import -import traceback from abc import abstractmethod, ABCMeta import hashlib import sys @@ -18,12 +17,13 @@ from typing import List, Optional, Set from netlib import wsgi, odict from netlib.exceptions import HttpException from netlib.http import Headers, http1, cookies +from netlib.utils import clean_bin from . import controller, tnetstring, filt, script, version, flow_format_compat from .onboarding import app from .proxy.config import HostMatcher from .protocol.http_replay import RequestReplayThread from .exceptions import Kill, FlowReadException -from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES +from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES, TCPFlow from collections import defaultdict @@ -900,6 +900,17 @@ class FlowMaster(controller.ServerMaster): self.handle_response(f) if f.error: self.handle_error(f) + elif isinstance(f, TCPFlow): + messages = f.messages + f.messages = [] + f.reply = controller.DummyReply() + self.handle_tcp_open(f) + while messages: + f.messages.append(messages.pop(0)) + self.handle_tcp_message(f) + if f.error: + self.handle_tcp_error(f) + self.handle_tcp_close(f) else: raise NotImplementedError() @@ -1087,18 +1098,52 @@ class FlowMaster(controller.ServerMaster): self.add_event('"{}" reloaded.'.format(s.filename), 'info') return ok - def handle_tcp_message(self, m): - self.run_script_hook("tcp_message", m) - m.reply() + def handle_tcp_open(self, flow): + self.state.add_flow(flow) + self.run_script_hook("tcp_open", flow) + flow.reply() + + def handle_tcp_message(self, flow): + self.run_script_hook("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(clean_bin(message.content), "debug") + flow.reply() + + def handle_tcp_error(self, flow): + if self.stream: + self.stream.add(flow) + self.add_event("Error in TCP connection to {}: {}".format( + repr(flow.server_conn.address), + flow.error + ), "info") + self.run_script_hook("tcp_error", flow) + flow.reply() + + def handle_tcp_close(self, flow): + self.state.delete_flow(flow) + if self.stream: + self.stream.add(flow) + self.run_script_hook("tcp_close", flow) + flow.reply() def shutdown(self): super(FlowMaster, self).shutdown() # Add all flows that are still active if self.stream: - for i in self.state.flows: - if not i.response: - self.stream.add(i) + for flow in self.state.flows: + # FIXME: We actually need to keep track of which flows are still active. + if isinstance(flow, HTTPFlow) and not flow.response: + self.stream.add(flow) + if isinstance(flow, TCPFlow): + # (assuming mitmdump only, this must be still active) + self.stream.add(flow) self.stop_stream() self.unload_scripts() diff --git a/mitmproxy/models/__init__.py b/mitmproxy/models/__init__.py index df86eff42..3d9d9dae6 100644 --- a/mitmproxy/models/__init__.py +++ b/mitmproxy/models/__init__.py @@ -7,9 +7,11 @@ from .http import ( from netlib.http import decoded from .connections import ClientConnection, ServerConnection from .flow import Flow, Error +from .tcp import TCPFlow FLOW_TYPES = dict( - http=HTTPFlow + http=HTTPFlow, + tcp=TCPFlow, ) __all__ = [ @@ -18,5 +20,6 @@ __all__ = [ "make_connect_response", "expect_continue_response", "ClientConnection", "ServerConnection", "Flow", "Error", + "TCPFlow" "FLOW_TYPES" ] diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py index 594147ec8..1019c9fb8 100644 --- a/mitmproxy/models/flow.py +++ b/mitmproxy/models/flow.py @@ -40,6 +40,9 @@ class Error(stateobject.StateObject): def __str__(self): return self.msg + def __repr__(self): + return self.msg + @classmethod def from_state(cls, state): # the default implementation assumes an empty constructor. Override @@ -99,6 +102,12 @@ class Flow(stateobject.StateObject): self._backup = state.pop("backup") super(Flow, self).set_state(state) + @classmethod + def from_state(cls, state): + f = cls(None, None) + f.set_state(state) + return f + def copy(self): f = copy.copy(self) diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py index 77a809cf5..75ffbfd0a 100644 --- a/mitmproxy/models/http.py +++ b/mitmproxy/models/http.py @@ -191,12 +191,6 @@ class HTTPFlow(Flow): response=HTTPResponse ) - @classmethod - def from_state(cls, state): - f = cls(None, None) - f.set_state(state) - return f - def __repr__(self): s = "".format(len(self.messages)) diff --git a/mitmproxy/protocol/rawtcp.py b/mitmproxy/protocol/rawtcp.py index 7d18025e8..5f6fca753 100644 --- a/mitmproxy/protocol/rawtcp.py +++ b/mitmproxy/protocol/rawtcp.py @@ -9,29 +9,26 @@ from netlib.exceptions import TcpException from netlib.tcp import ssl_read_select from netlib.utils import clean_bin from ..exceptions import ProtocolException +from ..models import Error +from ..models.tcp import TCPFlow, TCPMessage + 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 - def __init__(self, ctx, logging=True): - self.logging = logging + def __init__(self, ctx, ignore=False): + self.ignore = ignore super(RawTCPLayer, self).__init__(ctx) def __call__(self): self.connect() + if not self.ignore: + flow = TCPFlow(self.client_conn, self.server_conn, self) + self.channel.ask("tcp_open", flow) + buf = memoryview(bytearray(self.chunk_size)) client = self.client_conn.connection @@ -51,38 +48,24 @@ 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. - return + break else: dst.shutdown(socket.SHUT_WR) if len(conns) == 0: - return + break continue - 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, - # hence the "weird" direction string. - if dst == server: - direction = "-> tcp -> {}".format(repr(self.server_conn.address)) - else: - direction = "<- tcp <- {}".format(repr(self.server_conn.address)) - data = clean_bin(tcp_message.message) - self.log( - "{}\r\n{}".format(direction, data), - "info" - ) + tcp_message = TCPMessage(dst == server, buf[:size].tobytes()) + if not self.ignore: + flow.messages.append(tcp_message) + self.channel.ask("tcp_message", flow) + dst.sendall(tcp_message.content) except (socket.error, TcpException, SSL.Error) as e: - six.reraise( - ProtocolException, - ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e))), - sys.exc_info()[2] - ) + if not self.ignore: + flow.error = Error("TCP connection closed unexpectedly: {}".format(repr(e))) + self.channel.tell("tcp_error", flow) + finally: + if not self.ignore: + self.channel.tell("tcp_close", flow) diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py index 9caae02a5..21478dfdc 100644 --- a/mitmproxy/proxy/root_context.py +++ b/mitmproxy/proxy/root_context.py @@ -65,7 +65,7 @@ class RootContext(object): else: ignore = self.config.check_ignore((client_hello.client_sni, 443)) if ignore: - return RawTCPLayer(top_layer, logging=False) + return RawTCPLayer(top_layer, ignore=True) # 2. Always insert a TLS layer, even if there's neither client nor server tls. # An inline script may upgrade from http to https, From acd51befbb89f28397f8e3e52ead5bfa11fdc93f Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 11 May 2016 11:13:57 -0600 Subject: [PATCH 2/6] 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): From d3c30d9005e42a68cb3f5a5440f30f01f100cbec Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Wed, 11 May 2016 16:34:18 -0600 Subject: [PATCH 3/6] fix tests, don't double-add error'd flows --- mitmproxy/flow.py | 2 -- test/mitmproxy/tservers.py | 4 +++- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 1d05d4bb9..77ed87470 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -1116,8 +1116,6 @@ class FlowMaster(controller.ServerMaster): flow.reply() def handle_tcp_error(self, flow): - if self.stream: - self.stream.add(flow) self.add_event("Error in TCP connection to {}: {}".format( repr(flow.server_conn.address), flow.error diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index c9d68cfd0..a1e9c7139 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -161,7 +161,9 @@ class HTTPProxyTest(ProxyTestBase): q = "get:'/p/%s'" % spec else: q = "get:'%s/p/%s'" % (self.server.urlbase, spec) - return p.request(q) + resp = p.request(q) + p.close() + return resp def app(self, page): if self.ssl: From d38989fe7e4aade0b7ea6c8c7d0de80c2b28ec81 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 12 May 2016 14:29:57 -0600 Subject: [PATCH 4/6] tests: finish connections properly --- test/mitmproxy/tservers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index a1e9c7139..57225ee1b 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -162,6 +162,7 @@ class HTTPProxyTest(ProxyTestBase): else: q = "get:'%s/p/%s'" % (self.server.urlbase, spec) resp = p.request(q) + p.finish() p.close() return resp From f96697646ca495ddc045ab18134b8f70052c86ff Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 19 May 2016 22:09:00 -0700 Subject: [PATCH 5/6] add FlowMaster.active_flows --- mitmproxy/dump.py | 2 +- mitmproxy/flow.py | 22 +++++++++++----------- test/mitmproxy/tservers.py | 5 +---- 3 files changed, 13 insertions(+), 16 deletions(-) diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index aae397cd4..f1eabdb8c 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -320,7 +320,6 @@ class DumpMaster(flow.FlowMaster): self.outfile.flush() def _process_flow(self, f): - self.state.delete_flow(f) if self.filt and not f.match(self.filt): return @@ -328,6 +327,7 @@ class DumpMaster(flow.FlowMaster): def handle_request(self, f): flow.FlowMaster.handle_request(self, f) + self.state.delete_flow(f) if f: f.reply() return f diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 77ed87470..51ef71933 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -653,8 +653,9 @@ class FlowMaster(controller.ServerMaster): if server: self.add_server(server) self.state = state - self.server_playback = None - self.client_playback = None + self.active_flows = set() # type: Set[Flow] + self.server_playback = None # type: Optional[ServerPlaybackState] + self.client_playback = None # type: Optional[ClientPlaybackState] self.kill_nonreplay = False self.scripts = [] # type: List[script.Script] self.pause_scripts = False @@ -1033,6 +1034,7 @@ class FlowMaster(controller.ServerMaster): return if f not in self.state.flows: # don't add again on replay self.state.add_flow(f) + self.active_flows.add(f) self.replacehooks.run(f) self.setheaders.run(f) self.process_new_request(f) @@ -1053,6 +1055,7 @@ class FlowMaster(controller.ServerMaster): return f def handle_response(self, f): + self.active_flows.discard(f) self.state.update_flow(f) self.replacehooks.run(f) self.setheaders.run(f) @@ -1099,7 +1102,9 @@ class FlowMaster(controller.ServerMaster): return ok def handle_tcp_open(self, flow): - self.state.add_flow(flow) + # TODO: This would break mitmproxy currently. + # self.state.add_flow(flow) + self.active_flows.add(flow) self.run_script_hook("tcp_open", flow) flow.reply() @@ -1124,7 +1129,7 @@ class FlowMaster(controller.ServerMaster): flow.reply() def handle_tcp_close(self, flow): - self.state.delete_flow(flow) + self.active_flows.discard(flow) if self.stream: self.stream.add(flow) self.run_script_hook("tcp_close", flow) @@ -1135,13 +1140,8 @@ class FlowMaster(controller.ServerMaster): # Add all flows that are still active if self.stream: - for flow in self.state.flows: - # FIXME: We actually need to keep track of which flows are still active. - if isinstance(flow, HTTPFlow) and not flow.response: - self.stream.add(flow) - if isinstance(flow, TCPFlow): - # (assuming mitmdump only, this must be still active) - self.stream.add(flow) + for flow in self.active_flows: + self.stream.add(flow) self.stop_stream() self.unload_scripts() diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 57225ee1b..c9d68cfd0 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -161,10 +161,7 @@ class HTTPProxyTest(ProxyTestBase): q = "get:'/p/%s'" % spec else: q = "get:'%s/p/%s'" % (self.server.urlbase, spec) - resp = p.request(q) - p.finish() - p.close() - return resp + return p.request(q) def app(self, page): if self.ssl: From a3946d2a2d61a79f8f973c35f8321a37df3b8575 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Fri, 20 May 2016 13:27:26 -0700 Subject: [PATCH 6/6] tests++ --- test/mitmproxy/test_flow.py | 8 ++++++++ test/mitmproxy/tutils.py | 26 ++++++++++++++++++++++++-- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index b9c6a2f64..5c1e6a688 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) @@ -1205,6 +1209,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/tutils.py b/test/mitmproxy/tutils.py index d51ac1857..270a3911a 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 @@ -11,7 +13,7 @@ from unittest.case import SkipTest 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 ) @@ -44,6 +46,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 @@ -51,7 +73,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()