diff options
| -rw-r--r-- | mitmproxy/addonmanager.py | 55 | ||||
| -rw-r--r-- | mitmproxy/addons/script.py | 6 | ||||
| -rw-r--r-- | mitmproxy/controller.py | 3 | ||||
| -rw-r--r-- | mitmproxy/exceptions.py | 3 | ||||
| -rw-r--r-- | mitmproxy/master.py | 8 | ||||
| -rw-r--r-- | test/mitmproxy/test_addonmanager.py | 4 | ||||
| -rw-r--r-- | test/mitmproxy/test_examples.py | 8 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 2 |
8 files changed, 46 insertions, 43 deletions
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 43e76510..097f87b7 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -1,5 +1,6 @@ from mitmproxy import exceptions from mitmproxy import eventsequence +from . import ctx import pprint @@ -31,31 +32,27 @@ class AddonManager: return i def configure_all(self, options, updated): - self.invoke_all_with_context("configure", options, updated) - - def startup(self, s): - """ - Run startup events on addon. - """ - self.invoke_with_context(s, "start", self.master.options) + self.trigger("configure", options, updated) def add(self, *addons): """ Add addons to the end of the chain, and run their startup events. """ self.chain.extend(addons) - for i in addons: - self.startup(i) + with self.master.handlecontext(): + for i in addons: + self.invoke_addon(i, "start", self.master.options) def remove(self, addon): """ Remove an addon from the chain, and run its done events. """ self.chain = [i for i in self.chain if i is not addon] - self.invoke_with_context(addon, "done") + with self.master.handlecontext(): + self.invoke_addon(addon, "done") def done(self): - self.invoke_all_with_context("done") + self.trigger("done") def __len__(self): return len(self.chain) @@ -63,16 +60,15 @@ class AddonManager: def __str__(self): return pprint.pformat([str(i) for i in self.chain]) - def invoke_with_context(self, addon, name, *args, **kwargs): - with self.master.handlecontext(): - self.invoke(addon, name, *args, **kwargs) - - def invoke_all_with_context(self, name, *args, **kwargs): - with self.master.handlecontext(): - for i in self.chain: - self.invoke(i, name, *args, **kwargs) - - def invoke(self, addon, name, *args, **kwargs): + def invoke_addon(self, addon, name, *args, **kwargs): + """ + Invoke an event on an addon. This method must run within an + established handler context. + """ + if not ctx.master: + raise exceptions.AddonError( + "invoke_addon called without a handler context." + ) if name not in eventsequence.Events: # prama: no cover raise NotImplementedError("Unknown event") func = getattr(addon, name, None) @@ -83,9 +79,14 @@ class AddonManager: ) func(*args, **kwargs) - def __call__(self, name, *args, **kwargs): - for i in self.chain: - try: - self.invoke(i, name, *args, **kwargs) - except exceptions.AddonHalt: - return + def trigger(self, name, *args, **kwargs): + """ + Establish a handler context and trigger an event across all addons + """ + with self.master.handlecontext(): + for i in self.chain: + try: + self.invoke_addon(i, name, *args, **kwargs) + except exceptions.AddonHalt: + return + diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index cfbe5284..4d893f1c 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -273,11 +273,11 @@ class ScriptLoader: ctx.master.addons.chain = ochain[:pos + 1] + ordered + ochain[pos + 1:] for s in newscripts: - ctx.master.addons.startup(s) + ctx.master.addons.invoke_addon(s, "start", options) if self.is_running: # If we're already running, we configure and tell the addon # we're up and running. - ctx.master.addons.invoke_with_context( + ctx.master.addons.invoke_addon( s, "configure", options, options.keys() ) - ctx.master.addons.invoke_with_context(s, "running") + ctx.master.addons.invoke_addon(s, "running") diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 868d5841..aa4dcbbc 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -66,8 +66,7 @@ def handler(f): with master.handlecontext(): ret = f(master, message) - if handling: - master.addons(f.__name__, message) + master.addons.trigger(f.__name__, message) # Reset the handled flag - it's common for us to feed the same object # through handlers repeatedly, so we don't want this to persist across diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py index 309b8189..9b6328ac 100644 --- a/mitmproxy/exceptions.py +++ b/mitmproxy/exceptions.py @@ -102,6 +102,9 @@ class AddonError(MitmproxyException): class AddonHalt(MitmproxyException): + """ + Raised by addons to signal that no further handlers should handle this event. + """ pass diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 69359de6..19d069bc 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -65,8 +65,7 @@ class Master: """ level: debug, info, warn, error """ - with self.handlecontext(): - self.addons("log", log.LogEntry(e, level)) + self.addons.trigger("log", log.LogEntry(e, level)) def start(self): self.should_exit.clear() @@ -86,9 +85,8 @@ class Master: def tick(self, timeout): if self.first_tick: self.first_tick = False - self.addons.invoke_all_with_context("running") - with self.handlecontext(): - self.addons("tick") + self.addons.trigger("running") + self.addons.trigger("tick") changed = False try: mtype, obj = self.event_queue.get(timeout=timeout) diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 3e5f71c6..ef34371f 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -30,6 +30,6 @@ 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") diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index f20e0c8c..56692364 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -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/tservers.py b/test/mitmproxy/tservers.py index c47411ee..0f34e37e 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -80,7 +80,7 @@ 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") + self.addons.trigger("running") def clear_log(self): self.tlog = [] |