mirror of
https://github.com/zhigang1992/mitmproxy.git
synced 2026-04-29 20:55:03 +08:00
scripts: keep scripts just after the ScriptLoader in addon chain
We need scripts to run _before_ filestreamer, so we can't just add them to the end of the chain. This patch also fixes an issue that could cause scripts to be initialised un-necessarily if only the order of scripts in options changed.
This commit is contained in:
@@ -28,20 +28,32 @@ class Addons(object):
|
|||||||
with self.master.handlecontext():
|
with self.master.handlecontext():
|
||||||
i.configure(options, updated)
|
i.configure(options, updated)
|
||||||
|
|
||||||
|
def startup(self, s):
|
||||||
|
"""
|
||||||
|
Run startup events on addon.
|
||||||
|
"""
|
||||||
|
self.invoke_with_context(s, "start")
|
||||||
|
self.invoke_with_context(
|
||||||
|
s,
|
||||||
|
"configure",
|
||||||
|
self.master.options,
|
||||||
|
self.master.options.keys()
|
||||||
|
)
|
||||||
|
|
||||||
def add(self, *addons):
|
def add(self, *addons):
|
||||||
|
"""
|
||||||
|
Add addons to the end of the chain, and run their startup events.
|
||||||
|
"""
|
||||||
if not addons:
|
if not addons:
|
||||||
raise ValueError("No addons specified.")
|
raise ValueError("No addons specified.")
|
||||||
self.chain.extend(addons)
|
self.chain.extend(addons)
|
||||||
for i in addons:
|
for i in addons:
|
||||||
self.invoke_with_context(i, "start")
|
self.startup(i)
|
||||||
self.invoke_with_context(
|
|
||||||
i,
|
|
||||||
"configure",
|
|
||||||
self.master.options,
|
|
||||||
self.master.options.keys()
|
|
||||||
)
|
|
||||||
|
|
||||||
def remove(self, addon):
|
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.chain = [i for i in self.chain if i is not addon]
|
||||||
self.invoke_with_context(addon, "done")
|
self.invoke_with_context(addon, "done")
|
||||||
|
|
||||||
|
|||||||
@@ -239,16 +239,35 @@ class ScriptLoader():
|
|||||||
ctx.log.info("Un-loading script: %s" % a.name)
|
ctx.log.info("Un-loading script: %s" % a.name)
|
||||||
ctx.master.addons.remove(a)
|
ctx.master.addons.remove(a)
|
||||||
|
|
||||||
|
# The machinations below are to ensure that:
|
||||||
|
# - Scripts remain in the same order
|
||||||
|
# - Scripts are listed directly after the script addon. This is
|
||||||
|
# needed to ensure that interactions with, for instance, flow
|
||||||
|
# serialization remains correct.
|
||||||
|
# - Scripts are not initialized un-necessarily. If only a
|
||||||
|
# script's order in the script list has changed, it should simply
|
||||||
|
# be moved.
|
||||||
|
|
||||||
current = {}
|
current = {}
|
||||||
for a in ctx.master.addons.chain[:]:
|
for a in ctx.master.addons.chain[:]:
|
||||||
if isinstance(a, Script):
|
if isinstance(a, Script):
|
||||||
current[a.name] = a
|
current[a.name] = a
|
||||||
ctx.master.addons.chain.remove(a)
|
ctx.master.addons.chain.remove(a)
|
||||||
|
|
||||||
|
ordered = []
|
||||||
|
newscripts = []
|
||||||
for s in options.scripts:
|
for s in options.scripts:
|
||||||
if s in current:
|
if s in current:
|
||||||
ctx.master.addons.chain.append(current[s])
|
ordered.append(current[s])
|
||||||
else:
|
else:
|
||||||
ctx.log.info("Loading script: %s" % s)
|
ctx.log.info("Loading script: %s" % s)
|
||||||
sc = Script(s)
|
sc = Script(s)
|
||||||
ctx.master.addons.add(sc)
|
ordered.append(sc)
|
||||||
|
newscripts.append(sc)
|
||||||
|
|
||||||
|
ochain = ctx.master.addons.chain
|
||||||
|
pos = ochain.index(self)
|
||||||
|
ctx.master.addons.chain = ochain[:pos+1] + ordered + ochain[pos+1:]
|
||||||
|
|
||||||
|
for s in newscripts:
|
||||||
|
ctx.master.addons.startup(s)
|
||||||
|
|||||||
@@ -237,12 +237,8 @@ class TestScriptLoader(mastertest.MasterTest):
|
|||||||
"%s %s" % (rec, "b"),
|
"%s %s" % (rec, "b"),
|
||||||
]
|
]
|
||||||
debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"]
|
debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"]
|
||||||
assert debug == [
|
# No events, only order has changed
|
||||||
('debug', 'c configure'),
|
assert debug == []
|
||||||
('debug', 'a configure'),
|
|
||||||
('debug', 'b configure'),
|
|
||||||
]
|
|
||||||
m.event_log[:] = []
|
|
||||||
|
|
||||||
o.scripts = [
|
o.scripts = [
|
||||||
"%s %s" % (rec, "x"),
|
"%s %s" % (rec, "x"),
|
||||||
|
|||||||
@@ -16,6 +16,6 @@ def test_simple():
|
|||||||
o = options.Options()
|
o = options.Options()
|
||||||
m = controller.Master(o)
|
m = controller.Master(o)
|
||||||
a = addons.Addons(m)
|
a = addons.Addons(m)
|
||||||
a.add(o, TAddon("one"))
|
a.add(TAddon("one"))
|
||||||
assert a.get("one")
|
assert a.get("one")
|
||||||
assert not a.get("two")
|
assert not a.get("two")
|
||||||
|
|||||||
Reference in New Issue
Block a user