diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 9d0e70699..ac5610738 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -140,6 +140,7 @@ class AddonManager: """ for i in self.chain: self.remove(i) + self.lookup = {} def get(self, name): """ diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index 51341f2fb..decd07596 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -1,7 +1,7 @@ +import asyncio import os import importlib.util import importlib.machinery -import time import sys import types import typing @@ -59,13 +59,15 @@ def script_error_handler(path, exc, msg="", tb=False): ctx.log.error(log_msg) +ReloadInterval = 1 + + class Script: """ An addon that manages a single script. """ - ReloadInterval = 2 - def __init__(self, path): + def __init__(self, path: str, reload: bool) -> None: self.name = "scriptmanager:" + path self.path = path self.fullpath = os.path.expanduser( @@ -73,45 +75,57 @@ class Script: ) self.ns = None - self.last_load = 0 - self.last_mtime = 0 if not os.path.isfile(self.fullpath): raise exceptions.OptionsError('No such script') + self.reloadtask = None + if reload: + self.reloadtask = asyncio.ensure_future(self.watcher()) + else: + self.loadscript() + + def done(self): + if self.reloadtask: + self.reloadtask.cancel() + @property def addons(self): return [self.ns] if self.ns else [] - def tick(self): - if time.time() - self.last_load > self.ReloadInterval: + def loadscript(self): + ctx.log.info("Loading script %s" % self.path) + if self.ns: + ctx.master.addons.remove(self.ns) + self.ns = None + with addonmanager.safecall(): + ns = load_script(self.fullpath) + ctx.master.addons.register(ns) + self.ns = ns + if self.ns: + # We're already running, so we have to explicitly register and + # configure the addon + ctx.master.addons.invoke_addon(self.ns, "running") + ctx.master.addons.invoke_addon( + self.ns, + "configure", + ctx.options.keys() + ) + + async def watcher(self): + last_mtime = 0 + while True: try: mtime = os.stat(self.fullpath).st_mtime except FileNotFoundError: + ctx.log.info("Removing script %s" % self.path) scripts = list(ctx.options.scripts) scripts.remove(self.path) ctx.options.update(scripts=scripts) return - - if mtime > self.last_mtime: - ctx.log.info("Loading script: %s" % self.path) - if self.ns: - ctx.master.addons.remove(self.ns) - self.ns = None - with addonmanager.safecall(): - ns = load_script(self.fullpath) - ctx.master.addons.register(ns) - self.ns = ns - if self.ns: - # We're already running, so we have to explicitly register and - # configure the addon - ctx.master.addons.invoke_addon(self.ns, "running") - ctx.master.addons.invoke_addon( - self.ns, - "configure", - ctx.options.keys() - ) - self.last_load = time.time() - self.last_mtime = mtime + if mtime > last_mtime: + self.loadscript() + last_mtime = mtime + await asyncio.sleep(ReloadInterval) class ScriptLoader: @@ -125,9 +139,7 @@ class ScriptLoader: def load(self, loader): loader.add_option( "scripts", typing.Sequence[str], [], - """ - Execute a script. - """ + "Execute a script." ) def running(self): @@ -141,12 +153,7 @@ class ScriptLoader: simulated. """ try: - s = Script(path) - l = addonmanager.Loader(ctx.master) - ctx.master.addons.invoke_addon(s, "load", l) - ctx.master.addons.invoke_addon(s, "configure", ctx.options.keys()) - # Script is loaded on the first tick - ctx.master.addons.invoke_addon(s, "tick") + s = Script(path, False) for f in flows: for evt, arg in eventsequence.iterate(f): ctx.master.addons.invoke_addon(s, evt, arg) @@ -161,7 +168,7 @@ class ScriptLoader: for a in self.addons[:]: if a.path not in ctx.options.scripts: - ctx.log.info("Un-loading script: %s" % a.name) + ctx.log.info("Un-loading script: %s" % a.path) ctx.master.addons.remove(a) self.addons.remove(a) @@ -181,7 +188,7 @@ class ScriptLoader: if s in current: ordered.append(current[s]) else: - sc = Script(s) + sc = Script(s, True) ordered.append(sc) newscripts.append(sc) diff --git a/mitmproxy/master.py b/mitmproxy/master.py index bbbd07d0e..8eb016008 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -94,19 +94,14 @@ class Master: exc = None try: loop() - except Exception as e: + except Exception as e: # pragma: no cover exc = traceback.format_exc() finally: - if not self.should_exit.is_set(): + if not self.should_exit.is_set(): # pragma: no cover self.shutdown() - pending = asyncio.Task.all_tasks() loop = asyncio.get_event_loop() - try: - loop.run_until_complete(asyncio.gather(*pending)) - except Exception as e: - # When we exit with an error, shutdown might not happen cleanly, - # and we can get exceptions here caused by pending Futures. - pass + for p in asyncio.Task.all_tasks(): + p.cancel() loop.close() if exc: # pragma: no cover @@ -122,6 +117,7 @@ class Master: self.run_loop(loop.run_forever) async def _shutdown(self): + self.should_exit.set() if self.server: self.server.shutdown() loop = asyncio.get_event_loop() diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 73e574561..e0e4fef9a 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -123,11 +123,7 @@ class context: """ Loads a script from path, and returns the enclosed addon. """ - sc = script.Script(path) - loader = addonmanager.Loader(self.master) - self.master.addons.invoke_addon(sc, "load", loader) - self.configure(sc) - self.master.addons.invoke_addon(sc, "tick") + sc = script.Script(path, False) return sc.addons[0] if sc.addons else None def invoke(self, addon, event, *args, **kwargs): diff --git a/test/mitmproxy/addons/test_onboarding.py b/test/mitmproxy/addons/test_onboarding.py index 0d99b1ffd..a942062f6 100644 --- a/test/mitmproxy/addons/test_onboarding.py +++ b/test/mitmproxy/addons/test_onboarding.py @@ -4,23 +4,21 @@ from mitmproxy.addons import onboarding from mitmproxy.test import taddons from .. import tservers -import asyncio -import tornado.platform.asyncio -asyncio.set_event_loop_policy(tornado.platform.asyncio.AnyThreadEventLoopPolicy()) - class TestApp(tservers.HTTPProxyTest): def addons(self): return [onboarding.Onboarding()] - def test_basic(self): + @pytest.mark.asyncio + async def test_basic(self): ob = onboarding.Onboarding() with taddons.context(ob) as tctx: tctx.configure(ob) assert self.app("/").status_code == 200 @pytest.mark.parametrize("ext", ["pem", "p12"]) - def test_cert(self, ext): + @pytest.mark.asyncio + async def test_cert(self, ext): ob = onboarding.Onboarding() with taddons.context(ob) as tctx: tctx.configure(ob) @@ -29,7 +27,8 @@ class TestApp(tservers.HTTPProxyTest): assert resp.content @pytest.mark.parametrize("ext", ["pem", "p12"]) - def test_head(self, ext): + @pytest.mark.asyncio + async def test_head(self, ext): ob = onboarding.Onboarding() with taddons.context(ob) as tctx: tctx.configure(ob) diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 96e198419..fad514aef 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -12,6 +12,10 @@ from mitmproxy.test import tflow from mitmproxy.test import tutils +# We want this to be speedy for testing +script.ReloadInterval = 0.1 + + @pytest.mark.asyncio async def test_load_script(): with taddons.context() as tctx: @@ -71,7 +75,7 @@ class TestScript: def test_notfound(self): with taddons.context(): with pytest.raises(exceptions.OptionsError): - script.Script("nonexistent") + script.Script("nonexistent", False) def test_quotes_around_filename(self): """ @@ -81,21 +85,23 @@ class TestScript: path = tutils.test_data.path("mitmproxy/data/addonscripts/recorder/recorder.py") s = script.Script( - '"{}"'.format(path) + '"{}"'.format(path), + False ) assert '"' not in s.fullpath - def test_simple(self): + @pytest.mark.asyncio + async def test_simple(self): with taddons.context() as tctx: sc = script.Script( tutils.test_data.path( "mitmproxy/data/addonscripts/recorder/recorder.py" - ) + ), + True, ) tctx.master.addons.add(sc) tctx.configure(sc) - sc.tick() - + await tctx.master.await_log("recorder running") rec = tctx.master.addons.get("recorder") assert rec.call_log[0][0:2] == ("recorder", "load") @@ -112,25 +118,24 @@ class TestScript: f = tmpdir.join("foo.py") f.ensure(file=True) f.write("\n") - sc = script.Script(str(f)) + sc = script.Script(str(f), True) tctx.configure(sc) - sc.tick() assert await tctx.master.await_log("Loading") - tctx.master.clear() - sc.last_load, sc.last_mtime = 0, 0 - sc.tick() + tctx.master.clear() + f.write("\n") assert await tctx.master.await_log("Loading") @pytest.mark.asyncio async def test_exception(self): with taddons.context() as tctx: sc = script.Script( - tutils.test_data.path("mitmproxy/data/addonscripts/error.py") + tutils.test_data.path("mitmproxy/data/addonscripts/error.py"), + True, ) tctx.master.addons.add(sc) + await tctx.master.await_log("error running") tctx.configure(sc) - sc.tick() f = tflow.tflow(resp=True) tctx.master.addons.trigger("request", f) @@ -138,16 +143,17 @@ class TestScript: assert await tctx.master.await_log("ValueError: Error!") assert await tctx.master.await_log("error.py") - def test_addon(self): + @pytest.mark.asyncio + async def test_addon(self): with taddons.context() as tctx: sc = script.Script( tutils.test_data.path( "mitmproxy/data/addonscripts/addon.py" - ) + ), + True ) tctx.master.addons.add(sc) - tctx.configure(sc) - sc.tick() + await tctx.master.await_log("addon running") assert sc.ns.event_log == [ 'scriptload', 'addonload', 'scriptconfigure', 'addonconfigure' ] @@ -184,7 +190,6 @@ class TestScriptLoader: debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'recorder load', 'recorder running', 'recorder configure', - 'recorder tick', 'recorder requestheaders', 'recorder request', 'recorder responseheaders', 'recorder response' ] @@ -224,17 +229,21 @@ class TestScriptLoader: scripts = ["one", "one"] ) - def test_script_deletion(self): + @pytest.mark.asyncio + async def test_script_deletion(self): tdir = tutils.test_data.path("mitmproxy/data/addonscripts/") with open(tdir + "/dummy.py", 'w') as f: f.write("\n") + with taddons.context() as tctx: sl = script.ScriptLoader() tctx.master.addons.add(sl) tctx.configure(sl, scripts=[tutils.test_data.path("mitmproxy/data/addonscripts/dummy.py")]) + await tctx.master.await_log("Loading") os.remove(tutils.test_data.path("mitmproxy/data/addonscripts/dummy.py")) - tctx.invoke(sl, "tick") + + await tctx.master.await_log("Removing") assert not tctx.options.scripts assert not sl.addons @@ -286,17 +295,14 @@ class TestScriptLoader: 'a load', 'a running', 'a configure', - 'a tick', 'b load', 'b running', 'b configure', - 'b tick', 'c load', 'c running', 'c configure', - 'c tick', ] tctx.master.clear() @@ -317,7 +323,7 @@ class TestScriptLoader: 'b configure', ] - tctx.master.logs = [] + tctx.master.clear() tctx.configure( sc, scripts = [ @@ -325,9 +331,7 @@ class TestScriptLoader: "%s/a.py" % rec, ] ) - tctx.master.addons.invoke_addon(sc, "tick") - await tctx.master.await_log("a tick") - + await tctx.master.await_log("Loading") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'c done', @@ -336,6 +340,4 @@ class TestScriptLoader: 'e load', 'e running', 'e configure', - 'e tick', - 'a tick', ] diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index 8c834d82e..c6b540d46 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -1,3 +1,4 @@ +from mitmproxy import ctx event_log = [] @@ -7,6 +8,7 @@ class Addon: return event_log def load(self, opts): + ctx.log.info("addon running") event_log.append("addonload") def configure(self, updated): diff --git a/test/mitmproxy/data/addonscripts/error.py b/test/mitmproxy/data/addonscripts/error.py index 4a3c370fe..2f0c1755c 100644 --- a/test/mitmproxy/data/addonscripts/error.py +++ b/test/mitmproxy/data/addonscripts/error.py @@ -1,6 +1,9 @@ -def mkerr(): - raise ValueError("Error!") +from mitmproxy import ctx + + +def running(): + ctx.log.info("error running") def request(flow): - mkerr() + raise ValueError("Error!") diff --git a/test/mitmproxy/data/addonscripts/shutdown.py b/test/mitmproxy/data/addonscripts/shutdown.py new file mode 100644 index 000000000..f4a8f55d5 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/shutdown.py @@ -0,0 +1,5 @@ +from mitmproxy import ctx + + +def running(): + ctx.master.shutdown() \ No newline at end of file diff --git a/test/mitmproxy/data/addonscripts/stream_modify.py b/test/mitmproxy/data/addonscripts/stream_modify.py index 4fbf45c26..4ebcc3e94 100644 --- a/test/mitmproxy/data/addonscripts/stream_modify.py +++ b/test/mitmproxy/data/addonscripts/stream_modify.py @@ -1,7 +1,13 @@ +from mitmproxy import ctx + def modify(chunks): for chunk in chunks: yield chunk.replace(b"foo", b"bar") +def running(): + ctx.log.info("stream_modify running") + + def responseheaders(flow): flow.response.stream = modify diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index bf24e28b5..d78361046 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -256,11 +256,15 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): resp = p.request("get:'http://foo':h':foo'='bar'") assert resp.status_code == 400 - def test_stream_modify(self): + @pytest.mark.asyncio + async def test_stream_modify(self): s = script.Script( - tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py") + tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py"), + False, ) self.set_addons(s) + await self.master.await_log("stream_modify running") + d = self.pathod('200:b"foo"') assert d.content == b"bar" @@ -564,7 +568,8 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): def test_tcp_stream_modify(self): s = script.Script( - tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py") + tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py"), + False, ) self.set_addons(s) self._tcpproxy_on() diff --git a/test/mitmproxy/tools/test_main.py b/test/mitmproxy/tools/test_main.py index 575442769..751dcbe03 100644 --- a/test/mitmproxy/tools/test_main.py +++ b/test/mitmproxy/tools/test_main.py @@ -1,19 +1,23 @@ -import pytest +import asyncio from mitmproxy.tools import main -from mitmproxy import ctx +from mitmproxy.test import tutils + +shutdown_script = tutils.test_data.path("mitmproxy/data/addonscripts/shutdown.py") -@pytest.mark.asyncio -async def test_mitmweb(event_loop): +def test_mitmweb(event_loop): + asyncio.set_event_loop(event_loop) main.mitmweb([ "--no-web-open-browser", + "-s", shutdown_script, "-q", "-p", "0", ]) - await ctx.master._shutdown() -@pytest.mark.asyncio -async def test_mitmdump(): - main.mitmdump(["-q", "-p", "0"]) - await ctx.master._shutdown() +def test_mitmdump(event_loop): + asyncio.set_event_loop(event_loop) + main.mitmdump([ + "-s", shutdown_script, + "-q", "-p", "0", + ]) diff --git a/test/mitmproxy/tools/web/test_static_viewer.py b/test/mitmproxy/tools/web/test_static_viewer.py index dfc45bc2f..c044dee8c 100644 --- a/test/mitmproxy/tools/web/test_static_viewer.py +++ b/test/mitmproxy/tools/web/test_static_viewer.py @@ -1,5 +1,6 @@ import json from unittest import mock +import pytest from mitmproxy.test import taddons from mitmproxy.test import tflow @@ -57,7 +58,8 @@ def test_save_flows_content(ctx, tmpdir): assert p.join('response/content/Auto.json').check(file=1) -def test_static_viewer(tmpdir): +@pytest.mark.asyncio +async def test_static_viewer(tmpdir): s = static_viewer.StaticViewer() rf = readfile.ReadFile() sa = save.Save()