diff options
| -rw-r--r-- | mitmproxy/console/__init__.py | 8 | ||||
| -rw-r--r-- | mitmproxy/controller.py | 8 | ||||
| -rw-r--r-- | mitmproxy/dump.py | 16 | ||||
| -rw-r--r-- | mitmproxy/flow.py | 48 | ||||
| -rw-r--r-- | mitmproxy/models/flow.py | 2 | ||||
| -rw-r--r-- | mitmproxy/web/__init__.py | 6 | ||||
| -rw-r--r-- | test/mitmproxy/test_controller.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_dump.py | 22 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 50 | ||||
| -rw-r--r-- | test/mitmproxy/test_script.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_server.py | 32 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 16 |
12 files changed, 104 insertions, 108 deletions
diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py index 9ce02e72..c3157292 100644 --- a/mitmproxy/console/__init__.py +++ b/mitmproxy/console/__init__.py @@ -730,28 +730,28 @@ class ConsoleMaster(flow.FlowMaster): # Handlers @controller.handler - def handle_error(self, f): + def error(self, f): f = flow.FlowMaster.handle_error(self, f) if f: self.process_flow(f) return f @controller.handler - def handle_request(self, f): + def request(self, f): f = flow.FlowMaster.handle_request(self, f) if f: self.process_flow(f) return f @controller.handler - def handle_response(self, f): + def response(self, f): f = flow.FlowMaster.handle_response(self, f) if f: self.process_flow(f) return f @controller.handler - def handle_script_change(self, script): + def script_change(self, script): if super(ConsoleMaster, self).handle_script_change(script): signals.status_message.send(message='"{}" reloaded.'.format(script.filename)) else: diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 5a354fd2..51d58723 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -66,8 +66,10 @@ class Master(object): mtype, obj = self.event_queue.get(timeout=timeout) if mtype not in Events: raise ControlError("Unknown event %s"%repr(mtype)) - handle_func = getattr(self, "handle_" + mtype) - if not handle_func.func_dict.get("handler"): + handle_func = getattr(self, mtype) + if not hasattr(handle_func, "func_dict"): + raise ControlError("Handler %s not a function"%mtype) + if not handle_func.func_dict.get("__handler"): raise ControlError( "Handler function %s is not decorated with controller.handler"%( handle_func @@ -187,7 +189,7 @@ def handler(f): if handling and not message.reply.acked and not message.reply.taken: message.reply() return ret - wrapper.func_dict["handler"] = True + wrapper.func_dict["__handler"] = True return wrapper diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index cbf4b3da..4443995a 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -10,7 +10,6 @@ from . import flow, filt, contentviews, controller from .exceptions import ContentViewException, FlowReadException, ScriptException - class DumpError(Exception): pass @@ -326,22 +325,25 @@ class DumpMaster(flow.FlowMaster): self.echo_flow(f) - def handle_request(self, f): - flow.FlowMaster.handle_request(self, f) + @controller.handler + def request(self, f): + flow.FlowMaster.request(self, f) self.state.delete_flow(f) if f: f.reply() return f - def handle_response(self, f): - flow.FlowMaster.handle_response(self, f) + @controller.handler + def response(self, f): + flow.FlowMaster.response(self, f) if f: f.reply() self._process_flow(f) return f - def handle_error(self, f): - flow.FlowMaster.handle_error(self, f) + @controller.handler + def error(self, f): + flow.FlowMaster.error(self, f) if f: self._process_flow(f) return f diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 716d9a17..407f0d7b 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -208,9 +208,9 @@ class ClientPlaybackState: master.replay_request(self.current) else: self.current.reply = controller.DummyReply() - master.handle_request(self.current) + master.request(self.current) if self.current.response: - master.handle_response(self.current) + master.response(self.current) class ServerPlaybackState: @@ -894,23 +894,23 @@ class FlowMaster(controller.Master): f.reply = controller.DummyReply() if f.request: - self.handle_request(f) + self.request(f) if f.response: - self.handle_responseheaders(f) - self.handle_response(f) + self.responseheaders(f) + self.response(f) if f.error: - self.handle_error(f) + self.error(f) elif isinstance(f, TCPFlow): messages = f.messages f.messages = [] f.reply = controller.DummyReply() - self.handle_tcp_open(f) + self.tcp_open(f) while messages: f.messages.append(messages.pop(0)) - self.handle_tcp_message(f) + self.tcp_message(f) if f.error: - self.handle_tcp_error(f) - self.handle_tcp_close(f) + self.tcp_error(f) + self.tcp_close(f) else: raise NotImplementedError() @@ -987,31 +987,31 @@ class FlowMaster(controller.Master): rt.join() @controller.handler - def handle_log(self, l): + def log(self, l): self.add_event(l.msg, l.level) @controller.handler - def handle_clientconnect(self, root_layer): + def clientconnect(self, root_layer): self.run_script_hook("clientconnect", root_layer) @controller.handler - def handle_clientdisconnect(self, root_layer): + def clientdisconnect(self, root_layer): self.run_script_hook("clientdisconnect", root_layer) @controller.handler - def handle_serverconnect(self, server_conn): + def serverconnect(self, server_conn): self.run_script_hook("serverconnect", server_conn) @controller.handler - def handle_serverdisconnect(self, server_conn): + def serverdisconnect(self, server_conn): self.run_script_hook("serverdisconnect", server_conn) @controller.handler - def handle_next_layer(self, top_layer): + def next_layer(self, top_layer): self.run_script_hook("next_layer", top_layer) @controller.handler - def handle_error(self, f): + def error(self, f): self.state.update_flow(f) self.run_script_hook("error", f) if self.client_playback: @@ -1019,7 +1019,7 @@ class FlowMaster(controller.Master): return f @controller.handler - def handle_request(self, f): + def request(self, f): if f.live: app = self.apps.get(f.request) if app: @@ -1042,7 +1042,7 @@ class FlowMaster(controller.Master): return f @controller.handler - def handle_responseheaders(self, f): + def responseheaders(self, f): try: if self.stream_large_bodies: self.stream_large_bodies.run(f, False) @@ -1053,7 +1053,7 @@ class FlowMaster(controller.Master): return f @controller.handler - def handle_response(self, f): + def response(self, f): self.active_flows.discard(f) self.state.update_flow(f) self.replacehooks.run(f) @@ -1101,14 +1101,14 @@ class FlowMaster(controller.Master): return ok @controller.handler - def handle_tcp_open(self, flow): + def tcp_open(self, flow): # TODO: This would break mitmproxy currently. # self.state.add_flow(flow) self.active_flows.add(flow) self.run_script_hook("tcp_open", flow) @controller.handler - def handle_tcp_message(self, flow): + def tcp_message(self, flow): self.run_script_hook("tcp_message", flow) message = flow.messages[-1] direction = "->" if message.from_client else "<-" @@ -1120,7 +1120,7 @@ class FlowMaster(controller.Master): self.add_event(clean_bin(message.content), "debug") @controller.handler - def handle_tcp_error(self, flow): + def tcp_error(self, flow): self.add_event("Error in TCP connection to {}: {}".format( repr(flow.server_conn.address), flow.error @@ -1128,7 +1128,7 @@ class FlowMaster(controller.Master): self.run_script_hook("tcp_error", flow) @controller.handler - def handle_tcp_close(self, flow): + def tcp_close(self, flow): self.active_flows.discard(flow) if self.stream: self.stream.add(flow) diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py index 1019c9fb..8797fcd8 100644 --- a/mitmproxy/models/flow.py +++ b/mitmproxy/models/flow.py @@ -152,7 +152,7 @@ class Flow(stateobject.StateObject): self.error = Error("Connection killed") self.intercepted = False self.reply(Kill) - master.handle_error(self) + master.error(self) def intercept(self, master): """ diff --git a/mitmproxy/web/__init__.py b/mitmproxy/web/__init__.py index 7fef6df1..f8102ed8 100644 --- a/mitmproxy/web/__init__.py +++ b/mitmproxy/web/__init__.py @@ -197,17 +197,17 @@ class WebMaster(flow.FlowMaster): f.reply.take() @controller.handler - def handle_request(self, f): + def request(self, f): super(WebMaster, self).handle_request(f) self._process_flow(f) @controller.handler - def handle_response(self, f): + def response(self, f): super(WebMaster, self).handle_response(f) self._process_flow(f) @controller.handler - def handle_error(self, f): + def error(self, f): super(WebMaster, self).handle_error(f) self._process_flow(f) diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index 4b8656da..83ad428e 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -18,7 +18,7 @@ class TestMaster(object): def test_simple(self): class DummyMaster(controller.Master): @controller.handler - def handle_log(self, _): + def log(self, _): m.should_exit.set() def tick(self, timeout): diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py index ad4cee53..7d625c34 100644 --- a/test/mitmproxy/test_dump.py +++ b/test/mitmproxy/test_dump.py @@ -64,14 +64,14 @@ class TestDumpMaster: f = tutils.tflow(req=netlib.tutils.treq(content=content)) l = Log("connect") l.reply = mock.MagicMock() - m.handle_log(l) - m.handle_clientconnect(f.client_conn) - m.handle_serverconnect(f.server_conn) - m.handle_request(f) + m.log(l) + m.clientconnect(f.client_conn) + m.serverconnect(f.server_conn) + m.request(f) if not f.error: f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=content)) - f = m.handle_response(f) - m.handle_clientdisconnect(f.client_conn) + f = m.response(f) + m.clientdisconnect(f.client_conn) return f def _dummy_cycle(self, n, filt, content, **options): @@ -95,8 +95,8 @@ class TestDumpMaster: o = dump.Options(flow_detail=1) m = dump.DumpMaster(None, o, outfile=cs) f = tutils.tflow(err=True) - m.handle_request(f) - assert m.handle_error(f) + m.request(f) + assert m.error(f) assert "error" in cs.getvalue() def test_missing_content(self): @@ -105,10 +105,10 @@ class TestDumpMaster: m = dump.DumpMaster(None, o, outfile=cs) f = tutils.tflow() f.request.content = None - m.handle_request(f) + m.request(f) f.response = HTTPResponse.wrap(netlib.tutils.tresp()) f.response.content = None - m.handle_response(f) + m.response(f) assert "content missing" in cs.getvalue() def test_replay(self): @@ -160,7 +160,7 @@ class TestDumpMaster: assert o.verbosity == 2 def test_filter(self): - assert not "GET" in self._dummy_cycle(1, "~u foo", "", verbosity=1) + assert "GET" not in self._dummy_cycle(1, "~u foo", "", verbosity=1) def test_app(self): o = dump.Options(app=True) diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index f8338dcb..c5e39966 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -53,7 +53,7 @@ class TestStickyCookieState: assert s.domain_match("www.google.com", ".google.com") assert s.domain_match("google.com", ".google.com") - def test_handle_response(self): + def test_response(self): c = "SSID=mooo; domain=.google.com, FOO=bar; Domain=.google.com; Path=/; "\ "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; " @@ -100,7 +100,7 @@ class TestStickyCookieState: assert len(s.jar[googlekey].keys()) == 1 assert s.jar[googlekey]["somecookie"].items()[0][1] == "newvalue" - def test_handle_request(self): + def test_request(self): s, f = self._response("SSID=mooo", "www.google.com") assert "cookie" not in f.request.headers s.handle_request(f) @@ -109,7 +109,7 @@ class TestStickyCookieState: class TestStickyAuthState: - def test_handle_response(self): + def test_response(self): s = flow.StickyAuthState(filt.parse(".*")) f = tutils.tflow(resp=True) f.request.headers["authorization"] = "foo" @@ -798,8 +798,8 @@ class TestFlowMaster: fm = flow.FlowMaster(None, s) fm.load_script(tutils.test_data.path("scripts/reqerr.py")) f = tutils.tflow() - fm.handle_clientconnect(f.client_conn) - assert fm.handle_request(f) + fm.clientconnect(f.client_conn) + assert fm.request(f) def test_script(self): s = flow.State() @@ -807,18 +807,18 @@ class TestFlowMaster: fm.load_script(tutils.test_data.path("scripts/all.py")) f = tutils.tflow(resp=True) - fm.handle_clientconnect(f.client_conn) + fm.clientconnect(f.client_conn) assert fm.scripts[0].ns["log"][-1] == "clientconnect" - fm.handle_serverconnect(f.server_conn) + fm.serverconnect(f.server_conn) assert fm.scripts[0].ns["log"][-1] == "serverconnect" - fm.handle_request(f) + fm.request(f) assert fm.scripts[0].ns["log"][-1] == "request" - fm.handle_response(f) + fm.response(f) assert fm.scripts[0].ns["log"][-1] == "response" # load second script fm.load_script(tutils.test_data.path("scripts/all.py")) assert len(fm.scripts) == 2 - fm.handle_clientdisconnect(f.server_conn) + fm.clientdisconnect(f.server_conn) assert fm.scripts[0].ns["log"][-1] == "clientdisconnect" assert fm.scripts[1].ns["log"][-1] == "clientdisconnect" @@ -828,7 +828,7 @@ class TestFlowMaster: fm.load_script(tutils.test_data.path("scripts/all.py")) f.error = tutils.terr() - fm.handle_error(f) + fm.error(f) assert fm.scripts[0].ns["log"][-1] == "error" def test_duplicate_flow(self): @@ -853,20 +853,20 @@ class TestFlowMaster: fm.anticache = True fm.anticomp = True f = tutils.tflow(req=None) - fm.handle_clientconnect(f.client_conn) + fm.clientconnect(f.client_conn) f.request = HTTPRequest.wrap(netlib.tutils.treq()) - fm.handle_request(f) + fm.request(f) assert s.flow_count() == 1 f.response = HTTPResponse.wrap(netlib.tutils.tresp()) - fm.handle_response(f) + fm.response(f) assert s.flow_count() == 1 - fm.handle_clientdisconnect(f.client_conn) + fm.clientdisconnect(f.client_conn) f.error = Error("msg") f.error.reply = controller.DummyReply() - fm.handle_error(f) + fm.error(f) fm.load_script(tutils.test_data.path("scripts/a.py")) fm.shutdown() @@ -895,7 +895,7 @@ class TestFlowMaster: assert fm.state.flow_count() f.error = Error("error") - fm.handle_error(f) + fm.error(f) def test_server_playback(self): s = flow.State() @@ -976,12 +976,12 @@ class TestFlowMaster: fm.set_stickycookie(".*") f = tutils.tflow(resp=True) f.response.headers["set-cookie"] = "foo=bar" - fm.handle_request(f) - fm.handle_response(f) + fm.request(f) + fm.response(f) assert fm.stickycookie_state.jar assert not "cookie" in f.request.headers f = f.copy() - fm.handle_request(f) + fm.request(f) assert f.request.headers["cookie"] == "foo=bar" def test_stickyauth(self): @@ -996,12 +996,12 @@ class TestFlowMaster: fm.set_stickyauth(".*") f = tutils.tflow(resp=True) f.request.headers["authorization"] = "foo" - fm.handle_request(f) + fm.request(f) f = tutils.tflow(resp=True) assert fm.stickyauth_state.hosts assert not "authorization" in f.request.headers - fm.handle_request(f) + fm.request(f) assert f.request.headers["authorization"] == "foo" def test_stream(self): @@ -1017,15 +1017,15 @@ class TestFlowMaster: f = tutils.tflow(resp=True) fm.start_stream(file(p, "ab"), None) - fm.handle_request(f) - fm.handle_response(f) + fm.request(f) + fm.response(f) fm.stop_stream() assert r()[0].response f = tutils.tflow() fm.start_stream(file(p, "ab"), None) - fm.handle_request(f) + fm.request(f) fm.shutdown() assert not r()[1].response diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py index f321d15c..dd6f51ae 100644 --- a/test/mitmproxy/test_script.py +++ b/test/mitmproxy/test_script.py @@ -7,7 +7,7 @@ def test_duplicate_flow(): fm = flow.FlowMaster(None, s) fm.load_script(tutils.test_data.path("scripts/duplicate_flow.py")) f = tutils.tflow() - fm.handle_request(f) + fm.request(f) assert fm.state.flow_count() == 2 assert not fm.state.view[0].request.is_replay assert fm.state.view[1].request.is_replay diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index f4e7452f..bb4949e1 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -191,8 +191,8 @@ class TcpMixin: assert i_cert == i2_cert == n_cert # Make sure that TCP messages are in the event log. - assert any("305" in m for m in self.master.log) - assert any("306" in m for m in self.master.log) + assert any("305" in m for m in self.master.tlog) + assert any("306" in m for m in self.master.tlog) class AppMixin: @@ -261,7 +261,7 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin): p = self.pathoc() assert p.request(req % self.server.urlbase) assert p.request(req % self.server2.urlbase) - assert switched(self.proxy.log) + assert switched(self.proxy.tlog) def test_blank_leading_line(self): p = self.pathoc() @@ -500,7 +500,7 @@ class TestHttps2Http(tservers.ReverseProxyTest): def test_sni(self): p = self.pathoc(ssl=True, sni="example.com") assert p.request("get:'/p/200'").status_code == 200 - assert all("Error in handle_sni" not in msg for msg in self.proxy.log) + assert all("Error in handle_sni" not in msg for msg in self.proxy.tlog) def test_http(self): p = self.pathoc(ssl=False) @@ -625,7 +625,7 @@ class MasterRedirectRequest(tservers.TestMaster): redirect_port = None # Set by TestRedirectRequest @controller.handler - def handle_request(self, f): + def request(self, f): if f.request.path == "/p/201": # This part should have no impact, but it should also not cause any exceptions. @@ -636,13 +636,13 @@ class MasterRedirectRequest(tservers.TestMaster): # This is the actual redirection. f.request.port = self.redirect_port - super(MasterRedirectRequest, self).handle_request(f) + super(MasterRedirectRequest, self).request(f) @controller.handler - def handle_response(self, f): + def response(self, f): f.response.content = str(f.client_conn.address.port) f.response.headers["server-conn-id"] = str(f.server_conn.source_address.port) - super(MasterRedirectRequest, self).handle_response(f) + super(MasterRedirectRequest, self).response(f) class TestRedirectRequest(tservers.HTTPProxyTest): @@ -693,7 +693,7 @@ class MasterStreamRequest(tservers.TestMaster): Enables the stream flag on the flow for all requests """ @controller.handler - def handle_responseheaders(self, f): + def responseheaders(self, f): f.response.stream = True @@ -742,7 +742,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): class MasterFakeResponse(tservers.TestMaster): @controller.handler - def handle_request(self, f): + def request(self, f): resp = HTTPResponse.wrap(netlib.tutils.tresp()) f.reply(resp) @@ -763,14 +763,14 @@ class TestServerConnect(tservers.HTTPProxyTest): def test_unnecessary_serverconnect(self): """A replayed/fake response with no_upstream_cert should not connect to an upstream server""" assert self.pathod("200").status_code == 200 - for msg in self.proxy.tmaster.log: + for msg in self.proxy.tmaster.tlog: assert "serverconnect" not in msg class MasterKillRequest(tservers.TestMaster): @controller.handler - def handle_request(self, f): + def request(self, f): f.reply(Kill) @@ -787,7 +787,7 @@ class TestKillRequest(tservers.HTTPProxyTest): class MasterKillResponse(tservers.TestMaster): @controller.handler - def handle_response(self, f): + def response(self, f): f.reply(Kill) @@ -817,7 +817,7 @@ class TestTransparentResolveError(tservers.TransparentProxyTest): class MasterIncomplete(tservers.TestMaster): @controller.handler - def handle_request(self, f): + def request(self, f): resp = HTTPResponse.wrap(netlib.tutils.tresp()) resp.content = None f.reply(resp) @@ -949,14 +949,14 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest): kill_requests( self.chain[1].tmaster, - "handle_request", + "request", exclude = [ # fail first request 2, # allow second request ] ) - kill_requests(self.chain[0].tmaster, "handle_request", + kill_requests(self.chain[0].tmaster, "request", exclude=[ 1, # CONNECT # fail first request diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 51f4109d..24ebb476 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -39,19 +39,11 @@ class TestMaster(flow.FlowMaster): self.apps.add(errapp, "errapp", 80) self.clear_log() - @controller.handler - def handle_request(self, f): - flow.FlowMaster.handle_request(self, f) - - @controller.handler - def handle_response(self, f): - flow.FlowMaster.handle_response(self, f) - def clear_log(self): - self.log = [] + self.tlog = [] def add_event(self, message, level=None): - self.log.append(message) + self.tlog.append(message) class ProxyThread(threading.Thread): @@ -68,8 +60,8 @@ class ProxyThread(threading.Thread): return self.tmaster.server.address.port @property - def log(self): - return self.tmaster.log + def tlog(self): + return self.tmaster.tlog def run(self): self.tmaster.run() |