mitmdump: don't fail for other flowtypes

This commit is contained in:
Maximilian Hils 2016-04-29 20:18:32 -07:00
parent 60bbc250d4
commit 317b4becca
3 changed files with 26 additions and 18 deletions

View File

@ -21,7 +21,7 @@ from .onboarding import app
from .proxy.config import HostMatcher from .proxy.config import HostMatcher
from .protocol.http_replay import RequestReplayThread from .protocol.http_replay import RequestReplayThread
from .exceptions import Kill, FlowReadException 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 from collections import defaultdict
@ -873,9 +873,9 @@ class FlowMaster(controller.ServerMaster):
def load_flow(self, f): def load_flow(self, f):
""" """
Loads a flow, and returns a new flow object. Loads a flow
""" """
if isinstance(f, HTTPFlow):
if self.server and self.server.config.mode == "reverse": if self.server and self.server.config.mode == "reverse":
f.request.host = self.server.config.upstream_server.address.host f.request.host = self.server.config.upstream_server.address.host
f.request.port = self.server.config.upstream_server.address.port f.request.port = self.server.config.upstream_server.address.port
@ -889,7 +889,8 @@ class FlowMaster(controller.ServerMaster):
self.handle_response(f) self.handle_response(f)
if f.error: if f.error:
self.handle_error(f) self.handle_error(f)
return f else:
raise NotImplementedError()
def load_flows(self, fr): def load_flows(self, fr):
""" """
@ -1166,7 +1167,9 @@ class FlowReader:
raise FlowReadException(str(e)) raise FlowReadException(str(e))
if can_tell: if can_tell:
off = self.fo.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: except ValueError:
# Error is due to EOF # Error is due to EOF
if can_tell and self.fo.tell() == off and self.fo.read() == '': if can_tell and self.fo.tell() == off and self.fo.read() == '':

View File

@ -8,10 +8,15 @@ from netlib.http import decoded
from .connections import ClientConnection, ServerConnection from .connections import ClientConnection, ServerConnection
from .flow import Flow, Error from .flow import Flow, Error
FLOW_TYPES = dict(
http=HTTPFlow
)
__all__ = [ __all__ = [
"HTTPFlow", "HTTPRequest", "HTTPResponse", "Headers", "decoded", "HTTPFlow", "HTTPRequest", "HTTPResponse", "Headers", "decoded",
"make_error_response", "make_connect_request", "make_error_response", "make_connect_request",
"make_connect_response", "expect_continue_response", "make_connect_response", "expect_continue_response",
"ClientConnection", "ServerConnection", "ClientConnection", "ServerConnection",
"Flow", "Error", "Flow", "Error",
"FLOW_TYPES"
] ]

View File

@ -818,7 +818,7 @@ class TestFlowMaster:
s = flow.State() s = flow.State()
fm = flow.FlowMaster(None, s) fm = flow.FlowMaster(None, s)
f = tutils.tflow(resp=True) f = tutils.tflow(resp=True)
f = fm.load_flow(f) fm.load_flow(f)
assert s.flow_count() == 1 assert s.flow_count() == 1
f2 = fm.duplicate_flow(f) f2 = fm.duplicate_flow(f)
assert f2.response assert f2.response