diff options
author | Maximilian Hils <git@maximilianhils.com> | 2016-04-29 20:18:32 -0700 |
---|---|---|
committer | Maximilian Hils <git@maximilianhils.com> | 2016-04-29 20:18:32 -0700 |
commit | 317b4beccaa7e30f11095cfaf0364e325a875fd2 (patch) | |
tree | 92225b0134b47e8b1c9a02a7f1c420fb5281dd0c | |
parent | 60bbc250d45fbec544831fccee273f721a90435d (diff) | |
download | mitmproxy-317b4beccaa7e30f11095cfaf0364e325a875fd2.tar.gz mitmproxy-317b4beccaa7e30f11095cfaf0364e325a875fd2.tar.bz2 mitmproxy-317b4beccaa7e30f11095cfaf0364e325a875fd2.zip |
mitmdump: don't fail for other flowtypes
-rw-r--r-- | mitmproxy/flow.py | 39 | ||||
-rw-r--r-- | mitmproxy/models/__init__.py | 5 | ||||
-rw-r--r-- | test/mitmproxy/test_flow.py | 2 |
3 files changed, 27 insertions, 19 deletions
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 2bed4e1a..555c6895 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -21,7 +21,7 @@ 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 +from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES from collections import defaultdict @@ -873,23 +873,24 @@ class FlowMaster(controller.ServerMaster): def load_flow(self, f): """ - Loads a flow, and returns a new flow object. + Loads a flow """ - - if self.server and self.server.config.mode == "reverse": - f.request.host = self.server.config.upstream_server.address.host - f.request.port = self.server.config.upstream_server.address.port - f.request.scheme = self.server.config.upstream_server.scheme - - f.reply = controller.DummyReply() - if f.request: - self.handle_request(f) - if f.response: - self.handle_responseheaders(f) - self.handle_response(f) - if f.error: - self.handle_error(f) - return f + if isinstance(f, HTTPFlow): + if self.server and self.server.config.mode == "reverse": + f.request.host = self.server.config.upstream_server.address.host + f.request.port = self.server.config.upstream_server.address.port + f.request.scheme = self.server.config.upstream_server.scheme + + f.reply = controller.DummyReply() + if f.request: + self.handle_request(f) + if f.response: + self.handle_responseheaders(f) + self.handle_response(f) + if f.error: + self.handle_error(f) + else: + raise NotImplementedError() def load_flows(self, fr): """ @@ -1166,7 +1167,9 @@ class FlowReader: raise FlowReadException(str(e)) if can_tell: off = self.fo.tell() - yield HTTPFlow.from_state(data) + if data["type"] not in FLOW_TYPES: + raise FlowReadException("Unknown flow type: {}".format(data["type"])) + yield FLOW_TYPES[data["type"]].from_state(data) except ValueError: # Error is due to EOF if can_tell and self.fo.tell() == off and self.fo.read() == '': diff --git a/mitmproxy/models/__init__.py b/mitmproxy/models/__init__.py index f5f0213a..df86eff4 100644 --- a/mitmproxy/models/__init__.py +++ b/mitmproxy/models/__init__.py @@ -8,10 +8,15 @@ from netlib.http import decoded from .connections import ClientConnection, ServerConnection from .flow import Flow, Error +FLOW_TYPES = dict( + http=HTTPFlow +) + __all__ = [ "HTTPFlow", "HTTPRequest", "HTTPResponse", "Headers", "decoded", "make_error_response", "make_connect_request", "make_connect_response", "expect_continue_response", "ClientConnection", "ServerConnection", "Flow", "Error", + "FLOW_TYPES" ] diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 1d69a3f8..145e91cf 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -818,7 +818,7 @@ class TestFlowMaster: s = flow.State() fm = flow.FlowMaster(None, s) f = tutils.tflow(resp=True) - f = fm.load_flow(f) + fm.load_flow(f) assert s.flow_count() == 1 f2 = fm.duplicate_flow(f) assert f2.response |