Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-26 18:18:25 +00:00
Code Issues Projects Releases Wiki Activity

Clean up addonmanager interface

Browse Source 
Clarify the plethora of invocation methods we've sprouted, correct some usages
in the codebase.
This commit is contained in:
Aldo Cortesi 2017-03-16 07:53:19 +13:00 committed by Aldo Cortesi
parent ef582333ff
commit 169068c7ec
8 changed files with 46 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
mitmproxy/addonmanager.py
View File

@ -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

6
mitmproxy/addons/script.py
View File

@ -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")

3
mitmproxy/controller.py
View File

@ -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

3
mitmproxy/exceptions.py
View File

@ -102,6 +102,9 @@ class AddonError(MitmproxyException):
class AddonHalt(MitmproxyException):
"""
Raised by addons to signal that no further handlers should handle this event.
"""
pass

8
mitmproxy/master.py
View File

@ -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)

4
test/mitmproxy/test_addonmanager.py
View File

@ -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")

8
test/mitmproxy/test_examples.py
View File

@ -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:

2
test/mitmproxy/tservers.py
View File

@ -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 = []