script: add a namespace adaptor to match Addons

2 files changed, 26 insertions, 11 deletions

diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index ab068e47..468aa8e0 100644
--- a/mitmproxy/builtins/script.py
+++ b/mitmproxy/builtins/script.py
@@ -16,6 +16,19 @@ import watchdog.events
 from watchdog.observers import polling
 
 
+class NS:
+    def __init__(self, ns):
+        self.__dict__["ns"] = ns
+
+    def __getattr__(self, key):
+        if key not in self.ns:
+            raise AttributeError("No such element: %s", key)
+        return self.ns[key]
+
+    def __setattr__(self, key, value):
+        self.__dict__["ns"][key] = value
+
+
 def parse_command(command):
     """
         Returns a (path, args) tuple.
@@ -74,7 +87,7 @@ def load_script(path, args):
     ns = {'__file__': os.path.abspath(path)}
     with scriptenv(path, args):
         exec(code, ns, ns)
-    return ns
+    return NS(ns)
 
 
 class ReloadHandler(watchdog.events.FileSystemEventHandler):
@@ -118,27 +131,29 @@ class Script:
         # It's possible for ns to be un-initialised if we failed during
         # configure
         if self.ns is not None and not self.dead:
-            func = self.ns.get(name)
+            func = getattr(self.ns, name, None)
             if func:
                 with scriptenv(self.path, self.args):
-                    func(*args, **kwargs)
+                    return func(*args, **kwargs)
 
     def reload(self):
         self.should_reload.set()
 
+    def load_script(self):
+        self.ns = load_script(self.path, self.args)
+        self.run("start")
+
     def tick(self):
         if self.should_reload.is_set():
             self.should_reload.clear()
             ctx.log.info("Reloading script: %s" % self.name)
-            self.ns = load_script(self.path, self.args)
-            self.start()
+            self.load_script()
             self.configure(self.last_options)
         else:
             self.run("tick")
 
     def start(self):
-        self.ns = load_script(self.path, self.args)
-        self.run("start")
+        self.load_script()
 
     def configure(self, options):
         self.last_options = options
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index f37c7f94..1658088c 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -48,7 +48,7 @@ def test_load_script():
             "data/addonscripts/recorder.py"
         ), []
     )
-    assert ns["configure"]
+    assert ns.configure
 
 
 class TestScript(mastertest.MasterTest):
@@ -61,16 +61,16 @@ class TestScript(mastertest.MasterTest):
             )
         )
         m.addons.add(sc)
-        assert sc.ns["call_log"] == [
+        assert sc.ns.call_log == [
             ("solo", "start", (), {}),
             ("solo", "configure", (options.Options(),), {})
         ]
 
-        sc.ns["call_log"] = []
+        sc.ns.call_log = []
         f = tutils.tflow(resp=True)
         self.invoke(m, "request", f)
 
-        recf = sc.ns["call_log"][0]
+        recf = sc.ns.call_log[0]
         assert recf[1] == "request"
 
     def test_reload(self):