From c7d0850d8f697915b183f4fafd5ede7df2245569 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 15 Jul 2016 16:35:24 +1200 Subject: Script cleanups - Preserve script order on config change - Prohibit script duplicates (i.e. identical script + args) - Various cleanups and tweaks --- mitmproxy/builtins/script.py | 47 ++++++++++++++------- test/mitmproxy/builtins/test_script.py | 62 ++++++++++++++++++++++++++-- test/mitmproxy/data/addonscripts/recorder.py | 9 +++- 3 files changed, 100 insertions(+), 18 deletions(-) diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index bcb756f6..14f7dd5f 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -99,23 +99,25 @@ class Script: self.path, self.args = parse_command(command) self.ns = None self.observer = None + self.dead = False self.last_options = None self.should_reload = threading.Event() - for i in controller.Events - set(["start", "configure", "tick"]): - def mkprox(): - evt = i + for i in controller.Events: + if not hasattr(self, i): + def mkprox(): + evt = i - def prox(*args, **kwargs): - self.run(evt, *args, **kwargs) - return prox - setattr(self, i, mkprox()) + def prox(*args, **kwargs): + self.run(evt, *args, **kwargs) + return prox + setattr(self, i, mkprox()) def run(self, name, *args, **kwargs): # It's possible for ns to be un-initialised if we failed during # configure - if self.ns is not None: + if self.ns is not None and not self.dead: func = self.ns.get(name) if func: with scriptenv(self.path, self.args): @@ -149,18 +151,35 @@ class Script: self.observer.start() self.run("configure", options) + def done(self): + self.run("done") + self.dead = True + class ScriptLoader(): """ An addon that manages loading scripts from options. """ def configure(self, options): - for s in options.scripts or []: - if not ctx.master.addons.has_addon(s): + for s in options.scripts: + if options.scripts.count(s) > 1: + raise exceptions.OptionsError("Duplicate script: %s" % s) + + for a in ctx.master.addons.chain[:]: + if isinstance(a, Script) and a.name not in options.scripts: + ctx.log.info("Un-loading script: %s" % a.name) + ctx.master.addons.remove(a) + + current = {} + for a in ctx.master.addons.chain[:]: + if isinstance(a, Script): + current[a.name] = a + ctx.master.addons.chain.remove(a) + + for s in options.scripts: + if s in current: + ctx.master.addons.chain.append(current[s]) + else: ctx.log.info("Loading script: %s" % s) sc = Script(s) ctx.master.addons.add(sc) - for a in ctx.master.addons.chain: - if isinstance(a, Script): - if a.name not in options.scripts or []: - ctx.master.addons.remove(a) diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index 2c2568ed..da60b64c 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -58,8 +58,8 @@ class TestScript(mastertest.MasterTest): ) m.addons.add(sc) assert sc.ns["call_log"] == [ - ("start", (), {}), - ("configure", (options.Options(),), {}) + ("solo", "start", (), {}), + ("solo", "configure", (options.Options(),), {}) ] sc.ns["call_log"] = [] @@ -67,7 +67,7 @@ class TestScript(mastertest.MasterTest): self.invoke(m, "request", f) recf = sc.ns["call_log"][0] - assert recf[0] == "request" + assert recf[1] == "request" def test_reload(self): s = state.State() @@ -129,3 +129,59 @@ class TestScriptLoader(mastertest.MasterTest): assert len(m.addons) == 2 o.update(scripts = []) assert len(m.addons) == 1 + + def test_dupes(self): + s = state.State() + o = options.Options(scripts=["one", "one"]) + m = master.FlowMaster(o, None, s) + sc = script.ScriptLoader() + tutils.raises(exceptions.OptionsError, m.addons.add, sc) + + def test_order(self): + rec = tutils.test_data.path("data/addonscripts/recorder.py") + + s = state.State() + o = options.Options( + scripts = [ + "%s %s" % (rec, "a"), + "%s %s" % (rec, "b"), + "%s %s" % (rec, "c"), + ] + ) + m = mastertest.RecordingMaster(o, None, s) + sc = script.ScriptLoader() + m.addons.add(sc) + + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'a start'), ('debug', 'a configure'), + ('debug', 'b start'), ('debug', 'b configure'), + ('debug', 'c start'), ('debug', 'c configure') + ] + m.event_log[:] = [] + + o.scripts = [ + "%s %s" % (rec, "c"), + "%s %s" % (rec, "a"), + "%s %s" % (rec, "b"), + ] + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'c configure'), + ('debug', 'a configure'), + ('debug', 'b configure'), + ] + m.event_log[:] = [] + + o.scripts = [ + "%s %s" % (rec, "x"), + "%s %s" % (rec, "a"), + ] + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'c done'), + ('debug', 'b done'), + ('debug', 'x start'), + ('debug', 'x configure'), + ('debug', 'a configure'), + ] diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index 728203e3..b6ac8d89 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -1,8 +1,14 @@ from mitmproxy import controller from mitmproxy import ctx +import sys call_log = [] +if len(sys.argv) > 1: + name = sys.argv[1] +else: + name = "solo" + # Keep a log of all possible event calls evts = list(controller.Events) + ["configure"] for i in evts: @@ -10,9 +16,10 @@ for i in evts: evt = i def prox(*args, **kwargs): - lg = (evt, args, kwargs) + lg = (name, evt, args, kwargs) if evt != "log": ctx.log.info(str(lg)) call_log.append(lg) + ctx.log.debug("%s %s" % (name, evt)) return prox globals()[i] = mkprox() -- cgit v1.2.3