Merge pull request #2455 from mhils/script-fixes

Fix loading scripts with same filename

3 files changed, 38 insertions, 4 deletions

diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index b4274f8c..3e60fe67 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -1,5 +1,6 @@
 import os
-import importlib
+import importlib.util
+import importlib.machinery
 import time
 import sys
 import typing
@@ -16,12 +17,21 @@ def load_script(actx, path):
     if not os.path.exists(path):
         ctx.log.info("No such file: %s" % path)
         return
-    loader = importlib.machinery.SourceFileLoader(os.path.basename(path), path)
+
+    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)
     try:
         oldpath = sys.path
         sys.path.insert(0, os.path.dirname(path))
         with addonmanager.safecall():
-            m = loader.load_module()
+            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
@@ -64,7 +74,6 @@ class Script:
                 ctx.log.info("Loading script: %s" % self.path)
                 if self.ns:
                     ctx.master.addons.remove(self.ns)
-                    del sys.modules[self.ns.__name__]
                 self.ns = load_script(ctx, self.fullpath)
                 if self.ns:
                     # We're already running, so we have to explicitly register and
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 03b1f620..b7e6c82a 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -30,6 +30,30 @@ def test_load_script():
         assert not ns
 
 
+def test_load_fullname():
+    """
+    Test that loading two scripts at locations a/foo.py and b/foo.py works.
+    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"
+            )
+        )
+        assert ns.addons
+        ns2 = script.load_script(
+            tctx.ctx(),
+            tutils.test_data.path(
+                "mitmproxy/data/addonscripts/same_filename/addon.py"
+            )
+        )
+        assert ns.name != ns2.name
+        assert not hasattr(ns2, "addons")
+
+
 def test_script_print_stdout():
     with taddons.context() as tctx:
         with mock.patch('mitmproxy.ctx.log.warn') as mock_warn:
diff --git a/test/mitmproxy/data/addonscripts/same_filename/addon.py b/test/mitmproxy/data/addonscripts/same_filename/addon.py
new file mode 100644
index 00000000..c84a9b13
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/same_filename/addon.py
@@ -0,0 +1 @@
+foo = 42