From 704c1db1b7aa25c4ee6371bec811e92409b92398 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 24 Feb 2018 15:45:12 +1300 Subject: [PATCH] addon options: setheaders, stickyauth --- mitmproxy/addons/setheaders.py | 11 +++++++++++ mitmproxy/addons/stickyauth.py | 8 ++++++++ mitmproxy/options.py | 12 ------------ mitmproxy/test/taddons.py | 2 +- test/mitmproxy/addons/test_core.py | 19 +++++++------------ test/mitmproxy/addons/test_script.py | 11 +++++------ test/mitmproxy/addons/test_setheaders.py | 4 ++-- test/mitmproxy/addons/test_stickyauth.py | 4 ++-- 8 files changed, 36 insertions(+), 35 deletions(-) diff --git a/mitmproxy/addons/setheaders.py b/mitmproxy/addons/setheaders.py index d4d16e400..3517f70f3 100644 --- a/mitmproxy/addons/setheaders.py +++ b/mitmproxy/addons/setheaders.py @@ -1,3 +1,5 @@ +import typing + from mitmproxy import exceptions from mitmproxy import flowfilter from mitmproxy import ctx @@ -44,6 +46,15 @@ class SetHeaders: def __init__(self): self.lst = [] + def load(self, loader): + loader.add_option( + "setheaders", typing.Sequence[str], [], + """ + Header set pattern of the form "/pattern/header/value", where the + separator can be any character. + """ + ) + def configure(self, updated): if "setheaders" in updated: self.lst = [] diff --git a/mitmproxy/addons/stickyauth.py b/mitmproxy/addons/stickyauth.py index 24831d5ba..ab0599eef 100644 --- a/mitmproxy/addons/stickyauth.py +++ b/mitmproxy/addons/stickyauth.py @@ -1,3 +1,5 @@ +import typing + from mitmproxy import exceptions from mitmproxy import flowfilter from mitmproxy import ctx @@ -8,6 +10,12 @@ class StickyAuth: self.flt = None self.hosts = {} + def load(self, loader): + loader.add_option( + "stickyauth", typing.Optional[str], None, + "Set sticky auth filter. Matched against requests." + ) + def configure(self, updated): if "stickyauth" in updated: if ctx.options.stickyauth: diff --git a/mitmproxy/options.py b/mitmproxy/options.py index e42ff0e9f..bf2012d7a 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -67,7 +67,6 @@ class Options(optmanager.OptManager): view_filter = None # type: Optional[str] # FIXME: Options that should be uncomplicated to migrate to addons - setheaders = None # type: Sequence[str] stickyauth = None # type: Optional[str] stickycookie = None # type: Optional[str] stream_large_bodies = None # type: Optional[str] @@ -90,21 +89,10 @@ class Options(optmanager.OptManager): "showhost", bool, False, "Use the Host header to construct URLs for display." ) - self.add_option( - "setheaders", Sequence[str], [], - """ - Header set pattern of the form "/pattern/header/value", where the - separator can be any character. - """ - ) self.add_option( "stickycookie", Optional[str], None, "Set sticky cookie filter. Matched against requests." ) - self.add_option( - "stickyauth", Optional[str], None, - "Set sticky auth filter. Matched against requests." - ) self.add_option( "stream_large_bodies", Optional[str], None, """ diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 7efa2eaf8..5930e4141 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -68,7 +68,7 @@ class context: self.wrapped = None for a in addons: - self.master.addons.register(a) + self.master.addons.add(a) def ctx(self): """ diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py index 5c9a8a0d2..3f95bed4e 100644 --- a/test/mitmproxy/addons/test_core.py +++ b/test/mitmproxy/addons/test_core.py @@ -128,28 +128,23 @@ def test_options(tmpdir): p = str(tmpdir.join("path")) sa = core.Core() with taddons.context() as tctx: - tctx.options.stickycookie = "foo" - assert tctx.options.stickycookie == "foo" - sa.options_reset() - assert tctx.options.stickycookie is None - - tctx.options.stickycookie = "foo" - tctx.options.stickyauth = "bar" - sa.options_reset_one("stickycookie") - assert tctx.options.stickycookie is None - assert tctx.options.stickyauth == "bar" + tctx.options.listen_host = "foo" + assert tctx.options.listen_host == "foo" + sa.options_reset_one("listen_host") + assert tctx.options.listen_host != "foo" with pytest.raises(exceptions.CommandError): sa.options_reset_one("unknown") + tctx.options.listen_host = "foo" sa.options_save(p) with pytest.raises(exceptions.CommandError): sa.options_save("/") sa.options_reset() - assert tctx.options.stickyauth is None + assert tctx.options.listen_host == "" sa.options_load(p) - assert tctx.options.stickyauth == "bar" + assert tctx.options.listen_host == "foo" sa.options_load("/nonexistent") diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 78a5be6cf..73641d24a 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -171,7 +171,7 @@ class TestScriptLoader: "mitmproxy/data/addonscripts/recorder/recorder.py" ) sc = script.ScriptLoader() - with taddons.context() as tctx: + with taddons.context(sc) as tctx: sc.script_run([tflow.tflow(resp=True)], rp) debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ @@ -183,7 +183,7 @@ class TestScriptLoader: def test_script_run_nonexistent(self): sc = script.ScriptLoader() - with taddons.context(): + with taddons.context(sc): with pytest.raises(exceptions.CommandError): sc.script_run([tflow.tflow(resp=True)], "/") @@ -208,8 +208,7 @@ class TestScriptLoader: def test_dupes(self): sc = script.ScriptLoader() - with taddons.context() as tctx: - tctx.master.addons.add(sc) + with taddons.context(sc) as tctx: with pytest.raises(exceptions.OptionsError): tctx.configure( sc, @@ -232,7 +231,7 @@ class TestScriptLoader: def test_load_err(self): sc = script.ScriptLoader() - with taddons.context() as tctx: + with taddons.context(sc) as tctx: tctx.configure(sc, scripts=[ tutils.test_data.path("mitmproxy/data/addonscripts/load_error.py") ]) @@ -242,7 +241,7 @@ class TestScriptLoader: pass # this is expected and normally guarded. # on the next tick we should not fail however. tctx.invoke(sc, "tick") - assert len(tctx.master.addons) == 0 + assert len(tctx.master.addons) == 1 def test_order(self): rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder") diff --git a/test/mitmproxy/addons/test_setheaders.py b/test/mitmproxy/addons/test_setheaders.py index 3aaee7f40..2d5e9e3c8 100644 --- a/test/mitmproxy/addons/test_setheaders.py +++ b/test/mitmproxy/addons/test_setheaders.py @@ -19,14 +19,14 @@ class TestSetHeaders: def test_configure(self): sh = setheaders.SetHeaders() - with taddons.context() as tctx: + with taddons.context(sh) as tctx: with pytest.raises(Exception, match="Invalid setheader filter pattern"): tctx.configure(sh, setheaders = ["/~b/one/two"]) tctx.configure(sh, setheaders = ["/foo/bar/voing"]) def test_setheaders(self): sh = setheaders.SetHeaders() - with taddons.context() as tctx: + with taddons.context(sh) as tctx: tctx.configure( sh, setheaders = [ diff --git a/test/mitmproxy/addons/test_stickyauth.py b/test/mitmproxy/addons/test_stickyauth.py index ef7d07930..7b422fdd1 100644 --- a/test/mitmproxy/addons/test_stickyauth.py +++ b/test/mitmproxy/addons/test_stickyauth.py @@ -9,7 +9,7 @@ from mitmproxy import exceptions def test_configure(): r = stickyauth.StickyAuth() - with taddons.context() as tctx: + with taddons.context(r) as tctx: tctx.configure(r, stickyauth="~s") with pytest.raises(exceptions.OptionsError): tctx.configure(r, stickyauth="~~") @@ -20,7 +20,7 @@ def test_configure(): def test_simple(): r = stickyauth.StickyAuth() - with taddons.context() as tctx: + with taddons.context(r) as tctx: tctx.configure(r, stickyauth=".*") f = tflow.tflow(resp=True) f.request.headers["authorization"] = "foo"