diff options
Diffstat (limited to 'test')
27 files changed, 319 insertions, 239 deletions
diff --git a/test/mitmproxy/addons/test_check_alpn.py b/test/mitmproxy/addons/test_check_alpn.py index 2dc0c835..2b1d6058 100644 --- a/test/mitmproxy/addons/test_check_alpn.py +++ b/test/mitmproxy/addons/test_check_alpn.py @@ -12,7 +12,7 @@ class TestCheckALPN: with taddons.context() as tctx: a = check_alpn.CheckALPN() tctx.configure(a) - assert not any(msg in m for l, m in tctx.master.event_log) + assert not tctx.master.has_log(msg) def test_check_no_alpn(self, disable_alpn): msg = 'ALPN support missing' @@ -20,4 +20,4 @@ class TestCheckALPN: with taddons.context() as tctx: a = check_alpn.CheckALPN() tctx.configure(a) - assert any(msg in m for l, m in tctx.master.event_log) + assert tctx.master.has_log(msg) diff --git a/test/mitmproxy/addons/test_check_ca.py b/test/mitmproxy/addons/test_check_ca.py index fc64621c..cd34a9be 100644 --- a/test/mitmproxy/addons/test_check_ca.py +++ b/test/mitmproxy/addons/test_check_ca.py @@ -16,4 +16,4 @@ class TestCheckCA: tctx.master.server.config.certstore.default_ca.has_expired = mock.MagicMock(return_value=expired) a = check_ca.CheckCA() tctx.configure(a) - assert any(msg in m for l, m in tctx.master.event_log) is expired + assert tctx.master.has_log(msg) is expired diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py index c22b3589..f71662f0 100644 --- a/test/mitmproxy/addons/test_clientplayback.py +++ b/test/mitmproxy/addons/test_clientplayback.py @@ -23,7 +23,7 @@ class MockThread(): class TestClientPlayback: def test_playback(self): cp = clientplayback.ClientPlayback() - with taddons.context(): + with taddons.context() as tctx: assert cp.count() == 0 f = tflow.tflow(resp=True) cp.load([f]) @@ -35,17 +35,14 @@ class TestClientPlayback: assert rp.called assert cp.current_thread - cp.keepserving = False cp.flows = None cp.current_thread = None - with mock.patch("mitmproxy.master.Master.shutdown") as sd: - cp.tick() - assert sd.called + cp.tick() + assert tctx.master.has_event("processing_complete") cp.current_thread = MockThread() - with mock.patch("mitmproxy.master.Master.shutdown") as sd: - cp.tick() - assert cp.current_thread is None + cp.tick() + assert cp.current_thread is None def test_configure(self, tmpdir): cp = clientplayback.ClientPlayback() diff --git a/test/mitmproxy/addons/test_disable_h2c.py b/test/mitmproxy/addons/test_disable_h2c.py index d4df8390..cf20a368 100644 --- a/test/mitmproxy/addons/test_disable_h2c.py +++ b/test/mitmproxy/addons/test_disable_h2c.py @@ -31,7 +31,6 @@ class TestDisableH2CleartextUpgrade: b = io.BytesIO(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") f = tflow.tflow() f.request = http.HTTPRequest.wrap(http1.read_request(b)) - f.reply.handle() f.intercept() a.request(f) diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py index 47374617..fbcc4d16 100644 --- a/test/mitmproxy/addons/test_dumper.py +++ b/test/mitmproxy/addons/test_dumper.py @@ -151,7 +151,7 @@ class TestContentView: with taddons.context(options=options.Options()) as ctx: ctx.configure(d, flow_detail=4, verbosity=3) d.response(tflow.tflow()) - assert "Content viewer failed" in ctx.master.event_log[0][1] + assert ctx.master.has_log("content viewer failed") def test_tcp(): diff --git a/test/mitmproxy/addons/test_intercept.py b/test/mitmproxy/addons/test_intercept.py index 465e6433..f436a817 100644 --- a/test/mitmproxy/addons/test_intercept.py +++ b/test/mitmproxy/addons/test_intercept.py @@ -29,6 +29,5 @@ def test_simple(): assert not f.intercepted f = tflow.tflow(resp=True) - f.reply._state = "handled" r.response(f) assert f.intercepted diff --git a/test/mitmproxy/addons/test_keepserving.py b/test/mitmproxy/addons/test_keepserving.py new file mode 100644 index 00000000..70f7e1e6 --- /dev/null +++ b/test/mitmproxy/addons/test_keepserving.py @@ -0,0 +1,10 @@ +from mitmproxy.addons import keepserving +from mitmproxy.test import taddons + + +def test_keepserving(): + ks = keepserving.KeepServing() + + with taddons.context() as tctx: + ks.event_processing_complete() + assert tctx.master.should_exit.is_set() diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py new file mode 100644 index 00000000..b30c147b --- /dev/null +++ b/test/mitmproxy/addons/test_readfile.py @@ -0,0 +1,62 @@ +from mitmproxy.addons import readfile +from mitmproxy.test import taddons +from mitmproxy.test import tflow +from mitmproxy import io +from mitmproxy import exceptions +from unittest import mock + +import pytest + + +def write_data(path, corrupt=False): + with open(path, "wb") as tf: + w = io.FlowWriter(tf) + for i in range(3): + f = tflow.tflow(resp=True) + w.add(f) + for i in range(3): + f = tflow.tflow(err=True) + w.add(f) + f = tflow.ttcpflow() + w.add(f) + f = tflow.ttcpflow(err=True) + w.add(f) + if corrupt: + tf.write(b"flibble") + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_configure(mck, tmpdir): + + rf = readfile.ReadFile() + with taddons.context() as tctx: + tf = str(tmpdir.join("tfile")) + write_data(tf) + tctx.configure(rf, rfile=str(tf)) + assert not mck.called + rf.running() + assert mck.called + + write_data(tf, corrupt=True) + tctx.configure(rf, rfile=str(tf)) + with pytest.raises(exceptions.OptionsError): + rf.running() + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_corruption(mck, tmpdir): + + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_file("nonexistent") + assert not mck.called + assert len(tctx.master.logs) == 1 + + tfc = str(tmpdir.join("tfile")) + write_data(tfc, corrupt=True) + + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_file(tfc) + assert mck.called + assert len(tctx.master.logs) == 2 diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py new file mode 100644 index 00000000..76b01f4f --- /dev/null +++ b/test/mitmproxy/addons/test_readstdin.py @@ -0,0 +1,53 @@ + +import io +from mitmproxy.addons import readstdin +from mitmproxy.test import taddons +from mitmproxy.test import tflow +import mitmproxy.io +from unittest import mock + + +def gen_data(corrupt=False): + tf = io.BytesIO() + w = mitmproxy.io.FlowWriter(tf) + for i in range(3): + f = tflow.tflow(resp=True) + w.add(f) + for i in range(3): + f = tflow.tflow(err=True) + w.add(f) + f = tflow.ttcpflow() + w.add(f) + f = tflow.ttcpflow(err=True) + w.add(f) + if corrupt: + tf.write(b"flibble") + tf.seek(0) + return tf + + +class mStdin: + def __init__(self, d): + self.buffer = d + + def isatty(self): + return False + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_read(m, tmpdir): + rf = readstdin.ReadStdin() + with taddons.context() as tctx: + assert not m.called + rf.running(stdin=mStdin(gen_data())) + assert m.called + + rf.running(stdin=mStdin(None)) + assert tctx.master.logs + tctx.master.clear() + + m.reset_mock() + assert not m.called + rf.running(stdin=mStdin(gen_data(corrupt=True))) + assert m.called + assert tctx.master.logs diff --git a/test/mitmproxy/addons/test_replace.py b/test/mitmproxy/addons/test_replace.py index 7d590b35..9002afb5 100644 --- a/test/mitmproxy/addons/test_replace.py +++ b/test/mitmproxy/addons/test_replace.py @@ -97,6 +97,6 @@ class TestReplaceFile: tmpfile.remove() f = tflow.tflow() f.request.content = b"foo" - assert not tctx.master.event_log + assert not tctx.master.logs r.request(f) - assert tctx.master.event_log + assert tctx.master.logs diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 4c1b2e43..16827488 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -22,14 +22,12 @@ def test_scriptenv(): with taddons.context() as tctx: with script.scriptenv("path", []): raise SystemExit - assert tctx.master.event_log[0][0] == "error" - assert "exited" in tctx.master.event_log[0][1] + assert tctx.master.has_log("exited", "error") tctx.master.clear() with script.scriptenv("path", []): raise ValueError("fooo") - assert tctx.master.event_log[0][0] == "error" - assert "foo" in tctx.master.event_log[0][1] + assert tctx.master.has_log("fooo", "error") class Called: @@ -135,7 +133,7 @@ class TestScript: f.write(".") sc.tick() time.sleep(0.1) - if tctx.master.event_log: + if tctx.master.logs: return raise AssertionError("Change event not detected.") @@ -147,11 +145,11 @@ class TestScript: sc.start(tctx.options) f = tflow.tflow(resp=True) sc.request(f) - assert tctx.master.event_log[0][0] == "error" - assert len(tctx.master.event_log[0][1].splitlines()) == 6 - assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.event_log[0][1]) - assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.event_log[0][1]) - assert tctx.master.event_log[0][1].endswith("ValueError: Error!\n") + assert tctx.master.logs[0].level == "error" + assert len(tctx.master.logs[0].msg.splitlines()) == 6 + assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0].msg) + assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0].msg) + assert tctx.master.logs[0].msg.endswith("ValueError: Error!\n") def test_addon(self): with taddons.context() as tctx: @@ -256,21 +254,21 @@ class TestScriptLoader: "%s %s" % (rec, "c"), ] ) - debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"] + debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ - ('debug', 'a start'), - ('debug', 'a configure'), - ('debug', 'a running'), + 'a start', + 'a configure', + 'a running', - ('debug', 'b start'), - ('debug', 'b configure'), - ('debug', 'b running'), + 'b start', + 'b configure', + 'b running', - ('debug', 'c start'), - ('debug', 'c configure'), - ('debug', 'c running'), + 'c start', + 'c configure', + 'c running', ] - tctx.master.event_log = [] + tctx.master.logs = [] tctx.configure( sc, scripts = [ @@ -279,11 +277,10 @@ class TestScriptLoader: "%s %s" % (rec, "b"), ] ) - debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"] - # No events, only order has changed + debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [] - tctx.master.event_log = [] + tctx.master.logs = [] tctx.configure( sc, scripts = [ @@ -291,11 +288,11 @@ class TestScriptLoader: "%s %s" % (rec, "a"), ] ) - debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"] + debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ - ('debug', 'c done'), - ('debug', 'b done'), - ('debug', 'x start'), - ('debug', 'x configure'), - ('debug', 'x running'), + 'c done', + 'b done', + 'x start', + 'x configure', + 'x running', ] diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py index 54e4d281..02642c35 100644 --- a/test/mitmproxy/addons/test_serverplayback.py +++ b/test/mitmproxy/addons/test_serverplayback.py @@ -34,7 +34,7 @@ def test_tick(): s.final_flow = tflow.tflow() s.final_flow.live = False s.tick() - assert tctx.master.should_exit.is_set() + assert tctx.master.has_event("processing_complete") def test_server_playback(): @@ -315,7 +315,6 @@ def test_server_playback_full(): tctx.configure( s, refresh_server_playback = True, - keepserving=False ) f = tflow.tflow() diff --git a/test/mitmproxy/addons/test_termstatus.py b/test/mitmproxy/addons/test_termstatus.py index 01c14814..7becc857 100644 --- a/test/mitmproxy/addons/test_termstatus.py +++ b/test/mitmproxy/addons/test_termstatus.py @@ -6,7 +6,7 @@ def test_configure(): ts = termstatus.TermStatus() with taddons.context() as ctx: ts.running() - assert not ctx.master.event_log + assert not ctx.master.logs ctx.configure(ts, server=True) ts.running() - assert ctx.master.event_log + assert ctx.master.logs diff --git a/test/mitmproxy/contentviews/test_xml_html.py b/test/mitmproxy/contentviews/test_xml_html.py index 899ecfde..2b0aee4d 100644 --- a/test/mitmproxy/contentviews/test_xml_html.py +++ b/test/mitmproxy/contentviews/test_xml_html.py @@ -23,7 +23,7 @@ def test_format_xml(filename): path = data.path(filename) with open(path) as f: input = f.read() - with open(path.replace(".", "-formatted.")) as f: + with open("-formatted.".join(path.rsplit(".", 1))) as f: expected = f.read() tokens = xml_html.tokenize(input) assert xml_html.format_xml(tokens) == expected diff --git a/test/mitmproxy/net/http/test_request.py b/test/mitmproxy/net/http/test_request.py index 90ec31fe..ce49002c 100644 --- a/test/mitmproxy/net/http/test_request.py +++ b/test/mitmproxy/net/http/test_request.py @@ -1,7 +1,7 @@ from unittest import mock import pytest -from mitmproxy.net.http import Headers +from mitmproxy.net.http import Headers, Request from mitmproxy.test.tutils import treq from .test_message import _test_decoded_attr, _test_passthrough_attr @@ -35,6 +35,32 @@ class TestRequestCore: request.host = None assert repr(request) == "Request(GET /path)" + def test_make(self): + r = Request.make("GET", "https://example.com/") + assert r.method == "GET" + assert r.scheme == "https" + assert r.host == "example.com" + assert r.port == 443 + assert r.path == "/" + + r = Request.make("GET", "https://example.com/", "content", {"Foo": "bar"}) + assert r.content == b"content" + assert r.headers["content-length"] == "7" + assert r.headers["Foo"] == "bar" + + Request.make("GET", "https://example.com/", content=b"content") + with pytest.raises(TypeError): + Request.make("GET", "https://example.com/", content=42) + + r = Request.make("GET", "https://example.com/", headers=[(b"foo", b"bar")]) + assert r.headers["foo"] == "bar" + + r = Request.make("GET", "https://example.com/", headers=({"foo": "baz"})) + assert r.headers["foo"] == "baz" + + with pytest.raises(TypeError): + Request.make("GET", "https://example.com/", headers=42) + def test_replace(self): r = treq() r.path = b"foobarfoo" diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index 16efe415..447b15a5 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -5,7 +5,6 @@ import pytest from unittest import mock from mitmproxy.test import tutils -from mitmproxy import controller from mitmproxy import options from mitmproxy.addons import script from mitmproxy.addons import proxyauth @@ -250,17 +249,12 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): assert p.request(req) def test_get_connection_switching(self): - def switched(l): - for i in l: - if "serverdisconnect" in i: - return True - req = "get:'%s/p/200:b@1'" p = self.pathoc() with p.connect(): assert p.request(req % self.server.urlbase) assert p.request(req % self.server2.urlbase) - assert switched(self.proxy.tlog) + assert self.proxy.tmaster.has_log("serverdisconnect") def test_blank_leading_line(self): p = self.pathoc() @@ -602,7 +596,7 @@ class TestHttps2Http(tservers.ReverseProxyTest): p = self.pathoc(ssl=True, sni="example.com") with p.connect(): assert p.request("get:'/p/200'").status_code == 200 - assert all("Error in handle_sni" not in msg for msg in self.proxy.tlog) + assert not self.proxy.tmaster.has_log("error in handle_sni") def test_http(self): p = self.pathoc(ssl=False) @@ -731,13 +725,12 @@ class TestProxySSL(tservers.HTTPProxyTest): assert not first_flow.server_conn.via -class MasterRedirectRequest(tservers.TestMaster): - redirect_port = None # Set by TestRedirectRequest +class ARedirectRequest: + def __init__(self, redirect_port): + self.redirect_port = redirect_port - @controller.handler def request(self, f): if f.request.path == "/p/201": - # This part should have no impact, but it should also not cause any exceptions. addr = f.live.server_conn.address addr2 = ("127.0.0.1", self.redirect_port) @@ -746,17 +739,13 @@ class MasterRedirectRequest(tservers.TestMaster): # This is the actual redirection. f.request.port = self.redirect_port - super().request(f) - @controller.handler def response(self, f): f.response.content = bytes(f.client_conn.address[1]) f.response.headers["server-conn-id"] = str(f.server_conn.source_address[1]) - super().response(f) class TestRedirectRequest(tservers.HTTPProxyTest): - masterclass = MasterRedirectRequest ssl = True def test_redirect(self): @@ -769,7 +758,7 @@ class TestRedirectRequest(tservers.HTTPProxyTest): This test verifies that the original destination is restored for the third request. """ - self.master.redirect_port = self.server2.port + self.proxy.tmaster.addons.add(ARedirectRequest(self.server2.port)) p = self.pathoc() with p.connect(): @@ -778,13 +767,13 @@ class TestRedirectRequest(tservers.HTTPProxyTest): r1 = p.request("get:'/p/200'") assert r1.status_code == 200 assert self.server.last_log() - assert not self.server2.last_log() + assert not self.server2.expect_log(1, 0.5) self.server.clear_log() self.server2.clear_log() r2 = p.request("get:'/p/201'") assert r2.status_code == 201 - assert not self.server.last_log() + assert not self.server.expect_log(1, 0.5) assert self.server2.last_log() self.server.clear_log() @@ -792,25 +781,23 @@ class TestRedirectRequest(tservers.HTTPProxyTest): r3 = p.request("get:'/p/202'") assert r3.status_code == 202 assert self.server.last_log() - assert not self.server2.last_log() + assert not self.server2.expect_log(1, 0.5) assert r1.content == r2.content == r3.content -class MasterStreamRequest(tservers.TestMaster): +class AStreamRequest: """ Enables the stream flag on the flow for all requests """ - @controller.handler def responseheaders(self, f): f.response.stream = True class TestStreamRequest(tservers.HTTPProxyTest): - masterclass = MasterStreamRequest - def test_stream_simple(self): + self.proxy.tmaster.addons.add(AStreamRequest()) p = self.pathoc() with p.connect(): # a request with 100k of data but without content-length @@ -819,6 +806,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): assert len(r1.content) > 100000 def test_stream_multiple(self): + self.proxy.tmaster.addons.add(AStreamRequest()) p = self.pathoc() with p.connect(): # simple request with streaming turned on @@ -830,6 +818,7 @@ class TestStreamRequest(tservers.HTTPProxyTest): assert r1.status_code == 201 def test_stream_chunked(self): + self.proxy.tmaster.addons.add(AStreamRequest()) connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM) connection.connect(("127.0.0.1", self.proxy.port)) fconn = connection.makefile("rb") @@ -850,22 +839,20 @@ class TestStreamRequest(tservers.HTTPProxyTest): connection.close() -class MasterFakeResponse(tservers.TestMaster): - @controller.handler +class AFakeResponse: def request(self, f): f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) class TestFakeResponse(tservers.HTTPProxyTest): - masterclass = MasterFakeResponse def test_fake(self): + self.proxy.tmaster.addons.add(AFakeResponse()) f = self.pathod("200") assert "header-response" in f.headers class TestServerConnect(tservers.HTTPProxyTest): - masterclass = MasterFakeResponse ssl = True @classmethod @@ -876,39 +863,34 @@ class TestServerConnect(tservers.HTTPProxyTest): def test_unnecessary_serverconnect(self): """A replayed/fake response with no upstream_cert should not connect to an upstream server""" + self.proxy.tmaster.addons.add(AFakeResponse()) assert self.pathod("200").status_code == 200 - for msg in self.proxy.tmaster.tlog: - assert "serverconnect" not in msg + assert not self.proxy.tmaster.has_log("serverconnect") -class MasterKillRequest(tservers.TestMaster): +class AKillRequest: - @controller.handler def request(self, f): f.reply.kill() class TestKillRequest(tservers.HTTPProxyTest): - masterclass = MasterKillRequest - def test_kill(self): + self.proxy.tmaster.addons.add(AKillRequest()) with pytest.raises(exceptions.HttpReadDisconnect): self.pathod("200") # Nothing should have hit the server - assert not self.server.last_log() + assert not self.server.expect_log(1, 0.5) -class MasterKillResponse(tservers.TestMaster): - - @controller.handler +class AKillResponse: def response(self, f): f.reply.kill() class TestKillResponse(tservers.HTTPProxyTest): - masterclass = MasterKillResponse - def test_kill(self): + self.proxy.tmaster.addons.add(AKillResponse()) with pytest.raises(exceptions.HttpReadDisconnect): self.pathod("200") # The server should have seen a request @@ -922,9 +904,7 @@ class TestTransparentResolveError(tservers.TransparentProxyTest): assert self.pathod("304").status_code == 502 -class MasterIncomplete(tservers.TestMaster): - - @controller.handler +class AIncomplete: def request(self, f): resp = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) resp.content = None @@ -932,9 +912,8 @@ class MasterIncomplete(tservers.TestMaster): class TestIncompleteResponse(tservers.HTTPProxyTest): - masterclass = MasterIncomplete - def test_incomplete(self): + self.proxy.tmaster.addons.add(AIncomplete()) assert self.pathod("200").status_code == 502 diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index a9b6f0c4..0e397b8f 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -43,7 +43,7 @@ class TestConcurrent(tservers.MasterTest): ) ) sc.start(tctx.options) - assert "decorator not supported" in tctx.master.event_log[0][1] + assert tctx.master.has_log("decorator not supported") def test_concurrent_class(self): with taddons.context() as tctx: diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 3e5f71c6..e7be25b8 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -11,6 +11,7 @@ class TAddon: def __init__(self, name): self.name = name self.tick = True + self.custom_called = False def __repr__(self): return "Addon(%s)" % self.name @@ -18,11 +19,17 @@ class TAddon: def done(self): pass + def event_custom(self): + self.custom_called = True + def test_simple(): o = options.Options() m = master.Master(o, proxy.DummyServer(o)) a = addonmanager.AddonManager(m) + with pytest.raises(exceptions.AddonError): + a.invoke_addon(TAddon("one"), "done") + a.add(TAddon("one")) assert a.get("one") assert not a.get("two") @@ -30,6 +37,11 @@ def test_simple(): assert not a.chain a.add(TAddon("one")) - a("done") + a.trigger("done") with pytest.raises(exceptions.AddonError): - a("tick") + a.trigger("tick") + + ta = TAddon("one") + a.add(ta) + a.trigger("custom") + assert ta.custom_called diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index 6acd3e54..ccc8bf35 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -7,6 +7,7 @@ from mitmproxy.exceptions import Kill, ControlException from mitmproxy import controller from mitmproxy import master from mitmproxy import proxy +from mitmproxy.test import taddons class TMsg: @@ -15,22 +16,18 @@ class TMsg: class TestMaster: def test_simple(self): - class DummyMaster(master.Master): - @controller.handler + class tAddon: def log(self, _): - m.should_exit.set() + ctx.master.should_exit.set() - def tick(self, timeout): - # Speed up test - super().tick(0) - - m = DummyMaster(None, proxy.DummyServer(None)) - assert not m.should_exit.is_set() - msg = TMsg() - msg.reply = controller.DummyReply() - m.event_queue.put(("log", msg)) - m.run() - assert m.should_exit.is_set() + with taddons.context() as ctx: + ctx.master.addons.add(tAddon()) + assert not ctx.master.should_exit.is_set() + msg = TMsg() + msg.reply = controller.DummyReply() + ctx.master.event_queue.put(("log", msg)) + ctx.master.run() + assert ctx.master.should_exit.is_set() def test_server_simple(self): m = master.Master(None, proxy.DummyServer(None)) @@ -63,7 +60,6 @@ class TestChannel: def reply(): m, obj = q.get() assert m == "test" - obj.reply.handle() obj.reply.send(42) obj.reply.take() obj.reply.commit() @@ -85,10 +81,7 @@ class TestChannel: class TestReply: def test_simple(self): reply = controller.Reply(42) - assert reply.state == "unhandled" - - reply.handle() - assert reply.state == "handled" + assert reply.state == "start" reply.send("foo") assert reply.value == "foo" @@ -104,7 +97,6 @@ class TestReply: def test_kill(self): reply = controller.Reply(43) - reply.handle() reply.kill() reply.take() reply.commit() @@ -112,7 +104,6 @@ class TestReply: def test_ack(self): reply = controller.Reply(44) - reply.handle() reply.ack() reply.take() reply.commit() @@ -120,7 +111,6 @@ class TestReply: def test_reply_none(self): reply = controller.Reply(45) - reply.handle() reply.send(None) reply.take() reply.commit() @@ -128,7 +118,6 @@ class TestReply: def test_commit_no_reply(self): reply = controller.Reply(46) - reply.handle() reply.take() with pytest.raises(ControlException): reply.commit() @@ -137,7 +126,6 @@ class TestReply: def test_double_send(self): reply = controller.Reply(47) - reply.handle() reply.send(1) with pytest.raises(ControlException): reply.send(2) @@ -145,12 +133,11 @@ class TestReply: reply.commit() def test_state_transitions(self): - states = {"unhandled", "handled", "taken", "committed"} + states = {"start", "taken", "committed"} accept = { - "handle": {"unhandled"}, - "take": {"handled"}, + "take": {"start"}, "commit": {"taken"}, - "ack": {"handled", "taken"}, + "ack": {"start", "taken"}, } for fn, ok in accept.items(): for state in states: @@ -169,7 +156,6 @@ class TestReply: reply = controller.Reply(47) with pytest.raises(ControlException): reply.__del__() - reply.handle() reply.ack() reply.take() reply.commit() @@ -179,24 +165,22 @@ class TestDummyReply: def test_simple(self): reply = controller.DummyReply() for _ in range(2): - reply.handle() reply.ack() reply.take() reply.commit() reply.mark_reset() reply.reset() - assert reply.state == "unhandled" + assert reply.state == "start" def test_reset(self): reply = controller.DummyReply() - reply.handle() reply.ack() reply.take() reply.commit() reply.mark_reset() assert reply.state == "committed" reply.reset() - assert reply.state == "unhandled" + assert reply.state == "start" def test_del(self): reply = controller.DummyReply() diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index f20e0c8c..030f2c4e 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -41,7 +41,7 @@ class TestScripts(tservers.MasterTest): def test_add_header(self): m, _ = tscript("simple/add_header.py") f = tflow.tflow(resp=tutils.tresp()) - m.response(f) + m.addons.handle_lifecycle("response", f) assert f.response.headers["newheader"] == "foo" def test_custom_contentviews(self): @@ -56,7 +56,7 @@ class TestScripts(tservers.MasterTest): m, sc = tscript("simple/modify_body_inject_iframe.py", "http://example.org/evil_iframe") f = tflow.tflow(resp=tutils.tresp(content=b"<html><body>mitmproxy</body></html>")) - m.response(f) + m.addons.handle_lifecycle("response", f) content = f.response.content assert b'iframe' in content and b'evil_iframe' in content @@ -65,41 +65,41 @@ class TestScripts(tservers.MasterTest): form_header = Headers(content_type="application/x-www-form-urlencoded") f = tflow.tflow(req=tutils.treq(headers=form_header)) - m.request(f) + m.addons.handle_lifecycle("request", f) assert f.request.urlencoded_form["mitmproxy"] == "rocks" f.request.headers["content-type"] = "" - m.request(f) + m.addons.handle_lifecycle("request", f) assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] def test_modify_querystring(self): m, sc = tscript("simple/modify_querystring.py") f = tflow.tflow(req=tutils.treq(path="/search?q=term")) - m.request(f) + m.addons.handle_lifecycle("request", f) assert f.request.query["mitmproxy"] == "rocks" f.request.path = "/" - m.request(f) + m.addons.handle_lifecycle("request", f) assert f.request.query["mitmproxy"] == "rocks" def test_arguments(self): m, sc = tscript("simple/script_arguments.py", "mitmproxy rocks") f = tflow.tflow(resp=tutils.tresp(content=b"I <3 mitmproxy")) - m.response(f) + m.addons.handle_lifecycle("response", f) assert f.response.content == b"I <3 rocks" def test_redirect_requests(self): m, sc = tscript("simple/redirect_requests.py") f = tflow.tflow(req=tutils.treq(host="example.org")) - m.request(f) + m.addons.handle_lifecycle("request", f) assert f.request.host == "mitmproxy.org" def test_send_reply_from_proxy(self): m, sc = tscript("simple/send_reply_from_proxy.py") f = tflow.tflow(req=tutils.treq(host="example.com", port=80)) - m.request(f) + m.addons.handle_lifecycle("request", f) assert f.response.content == b"Hello World" def test_dns_spoofing(self): @@ -109,13 +109,13 @@ class TestScripts(tservers.MasterTest): host_header = Headers(host=original_host) f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) - m.requestheaders(f) + m.addons.handle_lifecycle("requestheaders", f) # Rewrite by reverse proxy mode f.request.scheme = "https" f.request.port = 443 - m.request(f) + m.addons.handle_lifecycle("request", f) assert f.request.scheme == "http" assert f.request.port == 80 @@ -145,7 +145,7 @@ class TestHARDump: path = str(tmpdir.join("somefile")) m, sc = tscript("complex/har_dump.py", shlex.quote(path)) - m.addons.invoke(m, "response", self.flow()) + m.addons.trigger("response", self.flow()) m.addons.remove(sc) with open(path, "r") as inp: @@ -156,7 +156,9 @@ class TestHARDump: path = str(tmpdir.join("somefile")) m, sc = tscript("complex/har_dump.py", shlex.quote(path)) - m.addons.invoke(m, "response", self.flow(resp_content=b"foo" + b"\xFF" * 10)) + m.addons.trigger( + "response", self.flow(resp_content=b"foo" + b"\xFF" * 10) + ) m.addons.remove(sc) with open(path, "r") as inp: @@ -194,7 +196,7 @@ class TestHARDump: path = str(tmpdir.join("somefile")) m, sc = tscript("complex/har_dump.py", shlex.quote(path)) - m.addons.invoke(m, "response", f) + m.addons.trigger("response", f) m.addons.remove(sc) with open(path, "r") as inp: diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index f4d32cbb..630fc7e4 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -3,12 +3,13 @@ import pytest from mitmproxy.test import tflow import mitmproxy.io -from mitmproxy import flowfilter, options +from mitmproxy import flowfilter +from mitmproxy import options +from mitmproxy.proxy import config from mitmproxy.contrib import tnetstring from mitmproxy.exceptions import FlowReadException from mitmproxy import flow from mitmproxy import http -from mitmproxy.proxy import ProxyConfig from mitmproxy.proxy.server import DummyServer from mitmproxy import master from . import tservers @@ -16,23 +17,6 @@ from . import tservers class TestSerialize: - def _treader(self): - sio = io.BytesIO() - w = mitmproxy.io.FlowWriter(sio) - for i in range(3): - f = tflow.tflow(resp=True) - w.add(f) - for i in range(3): - f = tflow.tflow(err=True) - w.add(f) - f = tflow.ttcpflow() - w.add(f) - f = tflow.ttcpflow(err=True) - w.add(f) - - sio.seek(0) - return mitmproxy.io.FlowReader(sio) - def test_roundtrip(self): sio = io.BytesIO() f = tflow.tflow() @@ -51,26 +35,6 @@ class TestSerialize: assert f2.request == f.request assert f2.marked - def test_load_flows(self): - r = self._treader() - s = tservers.TestState() - fm = master.Master(None, DummyServer()) - fm.addons.add(s) - fm.load_flows(r) - assert len(s.flows) == 6 - - def test_load_flows_reverse(self): - r = self._treader() - s = tservers.TestState() - opts = options.Options( - mode="reverse:https://use-this-domain" - ) - conf = ProxyConfig(opts) - fm = master.Master(opts, DummyServer(conf)) - fm.addons.add(s) - fm.load_flows(r) - assert s.flows[0].request.host == "use-this-domain" - def test_filter(self): sio = io.BytesIO() flt = flowfilter.parse("~c 200") @@ -122,6 +86,17 @@ class TestSerialize: class TestFlowMaster: + def test_load_flow_reverse(self): + s = tservers.TestState() + opts = options.Options( + mode="reverse:https://use-this-domain" + ) + conf = config.ProxyConfig(opts) + fm = master.Master(opts, DummyServer(conf)) + fm.addons.add(s) + f = tflow.tflow(resp=True) + fm.load_flow(f) + assert s.flows[0].request.host == "use-this-domain" def test_replay(self): fm = master.Master(None, DummyServer()) @@ -140,26 +115,26 @@ class TestFlowMaster: def test_create_flow(self): fm = master.Master(None, DummyServer()) - assert fm.create_request("GET", "http", "example.com", 80, "/") + assert fm.create_request("GET", "http://example.com/") def test_all(self): s = tservers.TestState() fm = master.Master(None, DummyServer()) fm.addons.add(s) f = tflow.tflow(req=None) - fm.clientconnect(f.client_conn) + fm.addons.handle_lifecycle("clientconnect", f.client_conn) f.request = http.HTTPRequest.wrap(mitmproxy.test.tutils.treq()) - fm.request(f) + fm.addons.handle_lifecycle("request", f) assert len(s.flows) == 1 f.response = http.HTTPResponse.wrap(mitmproxy.test.tutils.tresp()) - fm.response(f) + fm.addons.handle_lifecycle("response", f) assert len(s.flows) == 1 - fm.clientdisconnect(f.client_conn) + fm.addons.handle_lifecycle("clientdisconnect", f.client_conn) f.error = flow.Error("msg") - fm.error(f) + fm.addons.handle_lifecycle("error", f) fm.shutdown() diff --git a/test/mitmproxy/test_http.py b/test/mitmproxy/test_http.py index 889eb0a7..aa283530 100644 --- a/test/mitmproxy/test_http.py +++ b/test/mitmproxy/test_http.py @@ -175,7 +175,6 @@ class TestHTTPFlow: def test_kill(self): f = tflow.tflow() - f.reply.handle() f.intercept() assert f.killable f.kill() @@ -184,7 +183,6 @@ class TestHTTPFlow: def test_resume(self): f = tflow.tflow() - f.reply.handle() f.intercept() assert f.reply.state == "taken" f.resume() diff --git a/test/mitmproxy/test_taddons.py b/test/mitmproxy/test_taddons.py new file mode 100644 index 00000000..1e42141c --- /dev/null +++ b/test/mitmproxy/test_taddons.py @@ -0,0 +1,12 @@ + +from mitmproxy.test import taddons +from mitmproxy import ctx + + +def test_recordingmaster(): + with taddons.context() as tctx: + assert not tctx.master.has_log("nonexistent") + assert not tctx.master.has_event("nonexistent") + ctx.log.error("foo") + assert not tctx.master.has_log("foo", level="debug") + assert tctx.master.has_log("foo", level="error") diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 6c716ad1..44b9ff3f 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -5,6 +5,7 @@ from mitmproxy import proxy from mitmproxy import options from mitmproxy.tools.console import common from ... import tservers +import urwid def test_format_keyvals(): @@ -35,7 +36,10 @@ class TestMaster(tservers.MasterTest): def test_basic(self): m = self.mkmaster() for i in (1, 2, 3): - self.dummy_cycle(m, 1, b"") + try: + self.dummy_cycle(m, 1, b"") + except urwid.ExitMainLoop: + pass assert len(m.view) == i def test_run_script_once(self): @@ -48,11 +52,11 @@ class TestMaster(tservers.MasterTest): """regression test for https://github.com/mitmproxy/mitmproxy/issues/1605""" m = self.mkmaster(intercept="~b bar") f = tflow.tflow(req=tutils.treq(content=b"foo")) - m.request(f) + m.addons.handle_lifecycle("request", f) assert not m.view[0].intercepted f = tflow.tflow(req=tutils.treq(content=b"bar")) - m.request(f) + m.addons.handle_lifecycle("request", f) assert m.view[1].intercepted f = tflow.tflow(resp=tutils.tresp(content=b"bar")) - m.request(f) + m.addons.handle_lifecycle("request", f) assert m.view[2].intercepted diff --git a/test/mitmproxy/tools/test_dump.py b/test/mitmproxy/tools/test_dump.py index a15bf583..8e2fa5b2 100644 --- a/test/mitmproxy/tools/test_dump.py +++ b/test/mitmproxy/tools/test_dump.py @@ -2,7 +2,6 @@ import pytest from unittest import mock from mitmproxy import proxy -from mitmproxy import exceptions from mitmproxy import log from mitmproxy import controller from mitmproxy import options @@ -17,24 +16,12 @@ class TestDumpMaster(tservers.MasterTest): m = dump.DumpMaster(o, proxy.DummyServer(), with_termlog=False, with_dumper=False) return m - def test_read(self, tmpdir): - p = str(tmpdir.join("read")) - self.flowfile(p) - self.dummy_cycle( - self.mkmaster(None, rfile=p), - 1, b"", - ) - with pytest.raises(exceptions.OptionsError): - self.mkmaster(None, rfile="/nonexistent") - with pytest.raises(exceptions.OptionsError): - self.mkmaster(None, rfile="test_dump.py") - def test_has_error(self): m = self.mkmaster(None) ent = log.LogEntry("foo", "error") ent.reply = controller.DummyReply() - m.log(ent) - assert m.has_errored + m.addons.trigger("log", ent) + assert m.errorcheck.has_errored @pytest.mark.parametrize("termlog", [False, True]) def test_addons_termlog(self, termlog): diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py index 00dc2c7c..e3d5dc44 100644 --- a/test/mitmproxy/tools/web/test_app.py +++ b/test/mitmproxy/tools/web/test_app.py @@ -83,7 +83,6 @@ class TestApp(tornado.testing.AsyncHTTPTestCase): def test_resume(self): for f in self.view: - f.reply.handle() f.intercept() assert self.fetch( @@ -95,7 +94,6 @@ class TestApp(tornado.testing.AsyncHTTPTestCase): def test_kill(self): for f in self.view: f.backup() - f.reply.handle() f.intercept() assert self.fetch("/flows/42/kill", method="POST").code == 200 @@ -109,7 +107,6 @@ class TestApp(tornado.testing.AsyncHTTPTestCase): f = self.view.get_by_id("42") assert f - f.reply.handle() assert self.fetch("/flows/42", method="DELETE").code == 200 assert not self.view.get_by_id("42") diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index c47411ee..b737b82a 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -6,32 +6,27 @@ import sys import mitmproxy.platform from mitmproxy.proxy.config import ProxyConfig from mitmproxy.proxy.server import ProxyServer -from mitmproxy import master from mitmproxy import controller from mitmproxy import options from mitmproxy import exceptions from mitmproxy import io -from mitmproxy import http import pathod.test import pathod.pathoc +from mitmproxy import eventsequence from mitmproxy.test import tflow from mitmproxy.test import tutils +from mitmproxy.test import taddons class MasterTest: def cycle(self, master, content): f = tflow.tflow(req=tutils.treq(content=content)) - master.clientconnect(f.client_conn) - master.serverconnect(f.server_conn) - master.request(f) - if not f.error: - f.response = http.HTTPResponse.wrap( - tutils.tresp(content=content) - ) - master.response(f) - master.clientdisconnect(f) + master.addons.handle_lifecycle("clientconnect", f.client_conn) + for i in eventsequence.iterate(f): + master.addons.handle_lifecycle(*i) + master.addons.handle_lifecycle("clientdisconnect", f.client_conn) return f def dummy_cycle(self, master, n, content): @@ -68,11 +63,11 @@ class TestState: # self.flows.append(f) -class TestMaster(master.Master): +class TestMaster(taddons.RecordingMaster): def __init__(self, opts, config): s = ProxyServer(config) - master.Master.__init__(self, opts, s) + super().__init__(opts, s) def clear_addons(self, addons): self.addons.clear() @@ -80,18 +75,11 @@ class TestMaster(master.Master): self.addons.add(self.state) self.addons.add(*addons) self.addons.configure_all(self.options, self.options.keys()) - self.addons.invoke_all_with_context("running") - - def clear_log(self): - self.tlog = [] + self.addons.trigger("running") def reset(self, addons): self.clear_addons(addons) - self.clear_log() - - @controller.handler - def log(self, e): - self.tlog.append(e.msg) + self.clear() class ProxyThread(threading.Thread): @@ -111,7 +99,7 @@ class ProxyThread(threading.Thread): @property def tlog(self): - return self.tmaster.tlog + return self.tmaster.logs def run(self): self.tmaster.run() |