From 5a22496ee8f6f72bc5f75f623ca0d68d7a0d7855 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Tue, 9 Aug 2016 22:29:07 -0700 Subject: clean up code, improve DummyReply --- mitmproxy/console/flowlist.py | 4 ++-- mitmproxy/console/flowview.py | 3 +-- mitmproxy/controller.py | 23 +++++++++++++++++++---- mitmproxy/flow/master.py | 2 +- mitmproxy/flow/state.py | 2 +- mitmproxy/models/flow.py | 6 +++++- mitmproxy/web/app.py | 2 +- test/mitmproxy/builtins/test_anticache.py | 4 ++-- test/mitmproxy/builtins/test_anticomp.py | 4 ++-- test/mitmproxy/builtins/test_dumper.py | 2 +- test/mitmproxy/builtins/test_filestreamer.py | 6 +++--- test/mitmproxy/builtins/test_replace.py | 4 ++-- test/mitmproxy/builtins/test_script.py | 4 ++-- test/mitmproxy/builtins/test_setheaders.py | 8 ++++---- test/mitmproxy/builtins/test_stickyauth.py | 4 ++-- test/mitmproxy/builtins/test_stickycookie.py | 18 +++++++++--------- test/mitmproxy/mastertest.py | 17 +++++------------ test/mitmproxy/script/test_concurrent.py | 4 ++-- test/mitmproxy/test_controller.py | 5 ++++- test/mitmproxy/test_examples.py | 22 +++++++++++----------- test/mitmproxy/test_flow.py | 20 +++++++++++--------- 21 files changed, 90 insertions(+), 74 deletions(-) diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py index 571119ed..7e69e098 100644 --- a/mitmproxy/console/flowlist.py +++ b/mitmproxy/console/flowlist.py @@ -182,7 +182,7 @@ class ConnectionItem(urwid.WidgetWrap): self.flow.accept_intercept(self.master) signals.flowlist_change.send(self) elif key == "d": - if self.flow.reply and self.flow.reply.state != "committed": + if self.flow.killable: self.flow.kill(self.master) self.state.delete_flow(self.flow) signals.flowlist_change.send(self) @@ -246,7 +246,7 @@ class ConnectionItem(urwid.WidgetWrap): callback = self.save_flows_prompt, ) elif key == "X": - if self.flow.reply and self.flow.reply.state != "committed": + if self.flow.killable: self.flow.kill(self.master) elif key == "enter": if self.flow.request: diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py index 22314396..6d74be65 100644 --- a/mitmproxy/console/flowview.py +++ b/mitmproxy/console/flowview.py @@ -8,7 +8,6 @@ import urwid from typing import Optional, Union # noqa from mitmproxy import contentviews -from mitmproxy import controller from mitmproxy import models from mitmproxy import utils from mitmproxy.console import common @@ -537,7 +536,7 @@ class FlowView(tabs.Tabs): else: self.view_next_flow(self.flow) f = self.flow - if not f.reply.acked: + if f.killable: f.kill(self.master) self.state.delete_flow(f) elif key == "D": diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 5ef81f8c..72374f31 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -203,6 +203,12 @@ def handler(f): if not hasattr(message, "reply"): raise exceptions.ControlException("Message %s has no reply attribute" % message) + # DummyReplys may be reused multiple times. + # We only clear them up on the next handler so that we can access value and + # state in the meantime. + if isinstance(message.reply, DummyReply): + message.reply.reset() + # The following ensures that inheritance with wrapped handlers in the # base class works. If we're the first handler, then responsibility for # acking is ours. If not, it's someone else's and we ignore it. @@ -228,7 +234,7 @@ def handler(f): # DummyReplys may be reused multiple times. if isinstance(message.reply, DummyReply): - message.reply.reset() + message.reply.mark_reset() return ret # Mark this function as a handler wrapper wrapper.__dict__["__handler"] = True @@ -269,6 +275,10 @@ class Reply(object): def has_message(self): return self.value != NO_REPLY + @property + def done(self): + return self.state == "committed" + def handle(self): """ Reply are handled by controller.handlers, which may be nested. The first handler takes @@ -329,12 +339,17 @@ class DummyReply(Reply): """ def __init__(self): super(DummyReply, self).__init__(None) + self._should_reset = False - def reset(self): + def mark_reset(self): if self.state != "committed": raise exceptions.ControlException("Uncommitted reply: %s" % self.obj) - self._state = "unhandled" - self.value = NO_REPLY + self._should_reset = True + + def reset(self): + if self._should_reset: + self._state = "unhandled" + self.value = NO_REPLY def __del__(self): pass diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index e02cdc15..0475ef4e 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -234,7 +234,7 @@ class FlowMaster(controller.Master): pb = self.do_server_playback(f) if not pb and self.kill_nonreplay: self.add_log("Killed {}".format(f.request.url), "info") - f.kill(self) + f.reply.kill() def replay_request(self, f, block=False): """ diff --git a/mitmproxy/flow/state.py b/mitmproxy/flow/state.py index 6f7e6c03..8576fadc 100644 --- a/mitmproxy/flow/state.py +++ b/mitmproxy/flow/state.py @@ -178,7 +178,7 @@ class FlowStore(FlowList): def kill_all(self, master): for f in self._list: - if f.reply.state != "committed": + if f.killable: f.kill(master) diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py index 5d61e47f..fc673274 100644 --- a/mitmproxy/models/flow.py +++ b/mitmproxy/models/flow.py @@ -149,13 +149,17 @@ class Flow(stateobject.StateObject): self.set_state(self._backup) self._backup = None + @property + def killable(self): + return self.reply and self.reply.state in {"handled", "taken"} + def kill(self, master): """ Kill this request. """ self.error = Error("Connection killed") self.intercepted = False - # reply.state should only be handled or taken here. + # reply.state should only be "handled" or "taken" here. # if none of this is the case, .take() will raise an exception. if self.reply.state != "taken": self.reply.take() diff --git a/mitmproxy/web/app.py b/mitmproxy/web/app.py index 097de634..5bd6f274 100644 --- a/mitmproxy/web/app.py +++ b/mitmproxy/web/app.py @@ -234,7 +234,7 @@ class AcceptFlow(RequestHandler): class FlowHandler(RequestHandler): def delete(self, flow_id): - if self.flow.reply.state != "committed": + if self.flow.killable: self.flow.kill(self.master) self.state.delete_flow(self.flow) diff --git a/test/mitmproxy/builtins/test_anticache.py b/test/mitmproxy/builtins/test_anticache.py index ac321e26..8897de52 100644 --- a/test/mitmproxy/builtins/test_anticache.py +++ b/test/mitmproxy/builtins/test_anticache.py @@ -14,11 +14,11 @@ class TestAntiCache(mastertest.MasterTest): m.addons.add(o, sa) f = tutils.tflow(resp=True) - self.invoke(m, "request", f) + m.request(f) f = tutils.tflow(resp=True) f.request.headers["if-modified-since"] = "test" f.request.headers["if-none-match"] = "test" - self.invoke(m, "request", f) + m.request(f) assert "if-modified-since" not in f.request.headers assert "if-none-match" not in f.request.headers diff --git a/test/mitmproxy/builtins/test_anticomp.py b/test/mitmproxy/builtins/test_anticomp.py index a5f5a270..af9e4a6a 100644 --- a/test/mitmproxy/builtins/test_anticomp.py +++ b/test/mitmproxy/builtins/test_anticomp.py @@ -14,10 +14,10 @@ class TestAntiComp(mastertest.MasterTest): m.addons.add(o, sa) f = tutils.tflow(resp=True) - self.invoke(m, "request", f) + m.request(f) f = tutils.tflow(resp=True) f.request.headers["Accept-Encoding"] = "foobar" - self.invoke(m, "request", f) + m.request(f) assert f.request.headers["Accept-Encoding"] == "identity" diff --git a/test/mitmproxy/builtins/test_dumper.py b/test/mitmproxy/builtins/test_dumper.py index 1c7173e0..9b518b53 100644 --- a/test/mitmproxy/builtins/test_dumper.py +++ b/test/mitmproxy/builtins/test_dumper.py @@ -80,5 +80,5 @@ class TestContentView(mastertest.MasterTest): m = mastertest.RecordingMaster(o, None, s) d = dumper.Dumper() m.addons.add(o, d) - self.invoke(m, "response", tutils.tflow()) + m.response(tutils.tflow()) assert "Content viewer failed" in m.event_log[0][1] diff --git a/test/mitmproxy/builtins/test_filestreamer.py b/test/mitmproxy/builtins/test_filestreamer.py index 0e69b340..94d68813 100644 --- a/test/mitmproxy/builtins/test_filestreamer.py +++ b/test/mitmproxy/builtins/test_filestreamer.py @@ -28,8 +28,8 @@ class TestStream(mastertest.MasterTest): m.addons.add(o, sa) f = tutils.tflow(resp=True) - self.invoke(m, "request", f) - self.invoke(m, "response", f) + m.request(f) + m.response(f) m.addons.remove(sa) assert r()[0].response @@ -38,6 +38,6 @@ class TestStream(mastertest.MasterTest): m.addons.add(o, sa) f = tutils.tflow() - self.invoke(m, "request", f) + m.request(f) m.addons.remove(sa) assert not r()[1].response diff --git a/test/mitmproxy/builtins/test_replace.py b/test/mitmproxy/builtins/test_replace.py index 5e70ce56..07abcda4 100644 --- a/test/mitmproxy/builtins/test_replace.py +++ b/test/mitmproxy/builtins/test_replace.py @@ -43,10 +43,10 @@ class TestReplace(mastertest.MasterTest): f = tutils.tflow() f.request.content = b"foo" - self.invoke(m, "request", f) + m.request(f) assert f.request.content == b"bar" f = tutils.tflow(resp=True) f.response.content = b"foo" - self.invoke(m, "response", f) + m.response(f) assert f.response.content == b"bar" diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 2870fd17..0bac6ca0 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -69,7 +69,7 @@ class TestScript(mastertest.MasterTest): sc.ns.call_log = [] f = tutils.tflow(resp=True) - self.invoke(m, "request", f) + m.request(f) recf = sc.ns.call_log[0] assert recf[1] == "request" @@ -102,7 +102,7 @@ class TestScript(mastertest.MasterTest): ) m.addons.add(o, sc) f = tutils.tflow(resp=True) - self.invoke(m, "request", f) + m.request(f) assert m.event_log[0][0] == "error" def test_duplicate_flow(self): diff --git a/test/mitmproxy/builtins/test_setheaders.py b/test/mitmproxy/builtins/test_setheaders.py index 41c18360..63685177 100644 --- a/test/mitmproxy/builtins/test_setheaders.py +++ b/test/mitmproxy/builtins/test_setheaders.py @@ -33,12 +33,12 @@ class TestSetHeaders(mastertest.MasterTest): ) f = tutils.tflow() f.request.headers["one"] = "xxx" - self.invoke(m, "request", f) + m.request(f) assert f.request.headers["one"] == "two" f = tutils.tflow(resp=True) f.response.headers["one"] = "xxx" - self.invoke(m, "response", f) + m.response(f) assert f.response.headers["one"] == "three" m, sh = self.mkmaster( @@ -50,7 +50,7 @@ class TestSetHeaders(mastertest.MasterTest): f = tutils.tflow(resp=True) f.request.headers["one"] = "xxx" f.response.headers["one"] = "xxx" - self.invoke(m, "response", f) + m.response(f) assert f.response.headers.get_all("one") == ["two", "three"] m, sh = self.mkmaster( @@ -61,5 +61,5 @@ class TestSetHeaders(mastertest.MasterTest): ) f = tutils.tflow() f.request.headers["one"] = "xxx" - self.invoke(m, "request", f) + m.request(f) assert f.request.headers.get_all("one") == ["two", "three"] diff --git a/test/mitmproxy/builtins/test_stickyauth.py b/test/mitmproxy/builtins/test_stickyauth.py index 5757fb2d..00b12072 100644 --- a/test/mitmproxy/builtins/test_stickyauth.py +++ b/test/mitmproxy/builtins/test_stickyauth.py @@ -15,10 +15,10 @@ class TestStickyAuth(mastertest.MasterTest): f = tutils.tflow(resp=True) f.request.headers["authorization"] = "foo" - self.invoke(m, "request", f) + m.request(f) assert "address" in sa.hosts f = tutils.tflow(resp=True) - self.invoke(m, "request", f) + m.request(f) assert f.request.headers["authorization"] == "foo" diff --git a/test/mitmproxy/builtins/test_stickycookie.py b/test/mitmproxy/builtins/test_stickycookie.py index e9d92c83..13cea6a2 100644 --- a/test/mitmproxy/builtins/test_stickycookie.py +++ b/test/mitmproxy/builtins/test_stickycookie.py @@ -34,23 +34,23 @@ class TestStickyCookie(mastertest.MasterTest): f = tutils.tflow(resp=True) f.response.headers["set-cookie"] = "foo=bar" - self.invoke(m, "request", f) + m.request(f) f.reply.acked = False - self.invoke(m, "response", f) + m.response(f) assert sc.jar assert "cookie" not in f.request.headers f = f.copy() f.reply.acked = False - self.invoke(m, "request", f) + m.request(f) assert f.request.headers["cookie"] == "foo=bar" def _response(self, s, m, sc, cookie, host): f = tutils.tflow(req=ntutils.treq(host=host, port=80), resp=True) f.response.headers["Set-Cookie"] = cookie - self.invoke(m, "response", f) + m.response(f) return f def test_response(self): @@ -79,7 +79,7 @@ class TestStickyCookie(mastertest.MasterTest): c2 = "othercookie=helloworld; Path=/" f = self._response(s, m, sc, c1, "www.google.com") f.response.headers["Set-Cookie"] = c2 - self.invoke(m, "response", f) + m.response(f) googlekey = list(sc.jar.keys())[0] assert len(sc.jar[googlekey].keys()) == 2 @@ -96,7 +96,7 @@ class TestStickyCookie(mastertest.MasterTest): ] for c in cs: f.response.headers["Set-Cookie"] = c - self.invoke(m, "response", f) + m.response(f) googlekey = list(sc.jar.keys())[0] assert len(sc.jar[googlekey].keys()) == len(cs) @@ -108,7 +108,7 @@ class TestStickyCookie(mastertest.MasterTest): c2 = "somecookie=newvalue; Path=/" f = self._response(s, m, sc, c1, "www.google.com") f.response.headers["Set-Cookie"] = c2 - self.invoke(m, "response", f) + m.response(f) googlekey = list(sc.jar.keys())[0] assert len(sc.jar[googlekey].keys()) == 1 assert list(sc.jar[googlekey]["somecookie"].items())[0][1] == "newvalue" @@ -120,7 +120,7 @@ class TestStickyCookie(mastertest.MasterTest): # by setting the expire time in the past f = self._response(s, m, sc, "duffer=zafar; Path=/", "www.google.com") f.response.headers["Set-Cookie"] = "duffer=; Expires=Thu, 01-Jan-1970 00:00:00 GMT" - self.invoke(m, "response", f) + m.response(f) assert not sc.jar.keys() def test_request(self): @@ -128,5 +128,5 @@ class TestStickyCookie(mastertest.MasterTest): f = self._response(s, m, sc, "SSID=mooo", "www.google.com") assert "cookie" not in f.request.headers - self.invoke(m, "request", f) + m.request(f) assert "cookie" in f.request.headers diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 95597d2c..08659d19 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -1,5 +1,3 @@ -import mock - from . import tutils import netlib.tutils @@ -8,24 +6,19 @@ from mitmproxy import flow, proxy, models, controller class MasterTest: - def invoke(self, master, handler, *message): - with master.handlecontext(): - func = getattr(master, handler) - func(*message) def cycle(self, master, content): f = tutils.tflow(req=netlib.tutils.treq(content=content)) l = proxy.Log("connect") l.reply = controller.DummyReply() master.log(l) - self.invoke(master, "clientconnect", f.client_conn) - self.invoke(master, "clientconnect", f.client_conn) - self.invoke(master, "serverconnect", f.server_conn) - self.invoke(master, "request", f) + master.clientconnect(f.client_conn) + master.serverconnect(f.server_conn) + master.request(f) if not f.error: f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content)) - self.invoke(master, "response", f) - self.invoke(master, "clientdisconnect", f) + master.response(f) + master.clientdisconnect(f) return f def dummy_cycle(self, master, n, content): diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 07ba1c57..c4f1e9ae 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -25,8 +25,8 @@ class TestConcurrent(mastertest.MasterTest): ) m.addons.add(m.options, sc) f1, f2 = tutils.tflow(), tutils.tflow() - self.invoke(m, "request", f1) - self.invoke(m, "request", f2) + m.request(f1) + m.request(f2) start = time.time() while time.time() - start < 5: if f1.reply.state == f2.reply.state == "committed": diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index 8fe2453d..abf66b6c 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -187,6 +187,7 @@ class TestDummyReply(object): reply.ack() reply.take() reply.commit() + reply.mark_reset() reply.reset() assert reply.state == "unhandled" @@ -196,9 +197,11 @@ class TestDummyReply(object): reply.ack() reply.take() reply.commit() + reply.mark_reset() + assert reply.state == "committed" reply.reset() assert reply.state == "unhandled" def test_del(self): reply = controller.DummyReply() - reply.__del__() \ No newline at end of file + reply.__del__() diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index 34fcc261..6c24ace5 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -39,7 +39,7 @@ class TestScripts(mastertest.MasterTest): def test_add_header(self): m, _ = tscript("add_header.py") f = tutils.tflow(resp=netutils.tresp()) - self.invoke(m, "response", f) + m.response(f) assert f.response.headers["newheader"] == "foo" def test_custom_contentviews(self): @@ -54,9 +54,9 @@ class TestScripts(mastertest.MasterTest): tscript("iframe_injector.py") m, sc = tscript("iframe_injector.py", "http://example.org/evil_iframe") - flow = tutils.tflow(resp=netutils.tresp(content=b"mitmproxy")) - self.invoke(m, "response", flow) - content = flow.response.content + f = tutils.tflow(resp=netutils.tresp(content=b"mitmproxy")) + m.response(f) + content = f.response.content assert b'iframe' in content and b'evil_iframe' in content def test_modify_form(self): @@ -64,23 +64,23 @@ class TestScripts(mastertest.MasterTest): form_header = Headers(content_type="application/x-www-form-urlencoded") f = tutils.tflow(req=netutils.treq(headers=form_header)) - self.invoke(m, "request", f) + m.request(f) assert f.request.urlencoded_form[b"mitmproxy"] == b"rocks" f.request.headers["content-type"] = "" - self.invoke(m, "request", f) + m.request(f) assert list(f.request.urlencoded_form.items()) == [(b"foo", b"bar")] def test_modify_querystring(self): m, sc = tscript("modify_querystring.py") f = tutils.tflow(req=netutils.treq(path="/search?q=term")) - self.invoke(m, "request", f) + m.request(f) assert f.request.query["mitmproxy"] == "rocks" f.request.path = "/" - self.invoke(m, "request", f) + m.request(f) assert f.request.query["mitmproxy"] == "rocks" def test_modify_response_body(self): @@ -89,13 +89,13 @@ class TestScripts(mastertest.MasterTest): m, sc = tscript("modify_response_body.py", "mitmproxy rocks") f = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy")) - self.invoke(m, "response", f) + m.response(f) assert f.response.content == b"I <3 rocks" def test_redirect_requests(self): m, sc = tscript("redirect_requests.py") f = tutils.tflow(req=netutils.treq(host="example.org")) - self.invoke(m, "request", f) + m.request(f) assert f.request.host == "mitmproxy.org" def test_har_extractor(self): @@ -119,7 +119,7 @@ class TestScripts(mastertest.MasterTest): req=netutils.treq(**times), resp=netutils.tresp(**times) ) - self.invoke(m, "response", f) + m.response(f) m.addons.remove(sc) with open(path, "rb") as f: diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 256ee124..1caeb100 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -3,9 +3,9 @@ import io import netlib.utils from netlib.http import Headers -from mitmproxy import filt, controller, flow, options +from mitmproxy import filt, flow, options from mitmproxy.contrib import tnetstring -from mitmproxy.exceptions import FlowReadException +from mitmproxy.exceptions import FlowReadException, Kill from mitmproxy.models import Error from mitmproxy.models import Flow from mitmproxy.models import HTTPFlow @@ -372,14 +372,16 @@ class TestHTTPFlow(object): assert f.get_state() == f2.get_state() def test_kill(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) + fm = mock.Mock() f = tutils.tflow() f.reply.handle() - f.intercept(mock.Mock()) + f.intercept(fm) + assert fm.handle_intercept.called + assert f.killable f.kill(fm) - for i in s.view: - assert "killed" in str(i.error) + assert not f.killable + assert fm.error.called + assert f.reply.value == Kill def test_killall(self): s = flow.State() @@ -835,8 +837,8 @@ class TestFlowMaster: f = tutils.tflow() f.request.host = "nonexistent" - fm.process_new_request(f) - assert "killed" in f.error.msg + fm.request(f) + assert f.reply.value == Kill class TestRequest: -- cgit v1.2.3