diff options
| -rw-r--r-- | mitmproxy/addons/script.py | 32 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_script.py | 60 | ||||
| -rw-r--r-- | test/mitmproxy/test_addonmanager.py | 2 |
3 files changed, 43 insertions, 51 deletions
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index 3e60fe67..58e8cdcd 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -3,6 +3,7 @@ import importlib.util import importlib.machinery import time import sys +import types import typing from mitmproxy import addonmanager @@ -13,28 +14,23 @@ from mitmproxy import eventsequence from mitmproxy import ctx -def load_script(actx, path): - if not os.path.exists(path): - ctx.log.info("No such file: %s" % path) - return - +def load_script(path: str) -> types.ModuleType: fullname = "__mitmproxy_script__.{}".format( os.path.splitext(os.path.basename(path))[0] ) # the fullname is not unique among scripts, so if there already is an existing script with said # fullname, remove it. sys.modules.pop(fullname, None) + oldpath = sys.path + sys.path.insert(0, os.path.dirname(path)) try: - oldpath = sys.path - sys.path.insert(0, os.path.dirname(path)) - with addonmanager.safecall(): - loader = importlib.machinery.SourceFileLoader(fullname, path) - spec = importlib.util.spec_from_loader(fullname, loader=loader) - m = importlib.util.module_from_spec(spec) - loader.exec_module(m) - if not getattr(m, "name", None): - m.name = path - return m + loader = importlib.machinery.SourceFileLoader(fullname, path) + spec = importlib.util.spec_from_loader(fullname, loader=loader) + m = importlib.util.module_from_spec(spec) + loader.exec_module(m) + if not getattr(m, "name", None): + m.name = path # type: ignore + return m finally: sys.path[:] = oldpath @@ -65,7 +61,7 @@ class Script: try: mtime = os.stat(self.fullpath).st_mtime except FileNotFoundError: - scripts = ctx.options.scripts + scripts = list(ctx.options.scripts) scripts.remove(self.path) ctx.options.update(scripts=scripts) return @@ -74,7 +70,9 @@ class Script: ctx.log.info("Loading script: %s" % self.path) if self.ns: ctx.master.addons.remove(self.ns) - self.ns = load_script(ctx, self.fullpath) + self.ns = None + with addonmanager.safecall(): + self.ns = load_script(self.fullpath) if self.ns: # We're already running, so we have to explicitly register and # configure the addon diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index b7e6c82a..64fd9505 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -1,6 +1,5 @@ import traceback import sys -import time import os import pytest @@ -14,20 +13,17 @@ from mitmproxy.addons import script def test_load_script(): - with taddons.context() as tctx: - ns = script.load_script( - tctx.ctx(), - tutils.test_data.path( - "mitmproxy/data/addonscripts/recorder/recorder.py" - ) + ns = script.load_script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/recorder/recorder.py" ) - assert ns.addons + ) + assert ns.addons - ns = script.load_script( - tctx.ctx(), + with pytest.raises(FileNotFoundError): + script.load_script( "nonexistent" ) - assert not ns def test_load_fullname(): @@ -36,22 +32,19 @@ def test_load_fullname(): This only succeeds if they get assigned different basenames. """ - with taddons.context() as tctx: - ns = script.load_script( - tctx.ctx(), - tutils.test_data.path( - "mitmproxy/data/addonscripts/addon.py" - ) + ns = script.load_script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/addon.py" ) - assert ns.addons - ns2 = script.load_script( - tctx.ctx(), - tutils.test_data.path( - "mitmproxy/data/addonscripts/same_filename/addon.py" - ) + ) + assert ns.addons + ns2 = script.load_script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/same_filename/addon.py" ) - assert ns.name != ns2.name - assert not hasattr(ns2, "addons") + ) + assert ns.name != ns2.name + assert not hasattr(ns2, "addons") def test_script_print_stdout(): @@ -59,7 +52,6 @@ def test_script_print_stdout(): with mock.patch('mitmproxy.ctx.log.warn') as mock_warn: with addonmanager.safecall(): ns = script.load_script( - tctx.ctx(), tutils.test_data.path( "mitmproxy/data/addonscripts/print.py" ) @@ -103,11 +95,13 @@ class TestScript: sc = script.Script(str(f)) tctx.configure(sc) sc.tick() - for _ in range(3): - sc.last_load, sc.last_mtime = 0, 0 - sc.tick() - time.sleep(0.1) - tctx.master.has_log("Loading") + assert tctx.master.has_log("Loading") + tctx.master.clear() + assert not tctx.master.has_log("Loading") + + sc.last_load, sc.last_mtime = 0, 0 + sc.tick() + assert tctx.master.has_log("Loading") def test_exception(self): with taddons.context() as tctx: @@ -121,8 +115,8 @@ class TestScript: f = tflow.tflow(resp=True) tctx.master.addons.trigger("request", f) - tctx.master.has_log("ValueError: Error!") - tctx.master.has_log("error.py") + assert tctx.master.has_log("ValueError: Error!") + assert tctx.master.has_log("error.py") def test_addon(self): with taddons.context() as tctx: diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 5bff61d1..722d108b 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -115,7 +115,7 @@ def test_simple(): a.add(TAddon("one")) a.trigger("done") a.trigger("tick") - tctx.master.has_log("not callable") + assert tctx.master.has_log("not callable") tctx.master.clear() a.get("one").tick = addons |