Merge pull request #1410 from cortesi/addons

Keep maturing scripts and addons

7 files changed, 136 insertions, 70 deletions

diff --git a/examples/filt.py b/examples/filt.py
index 21744edd..9ccf9fa1 100644
--- a/examples/filt.py
+++ b/examples/filt.py
@@ -1,18 +1,20 @@
-# This scripts demonstrates how to use mitmproxy's filter pattern in inline scripts.
+# This scripts demonstrates how to use mitmproxy's filter pattern in scripts.
 # Usage: mitmdump -s "filt.py FILTER"
 import sys
 from mitmproxy import filt
 
-state = {}
+
+class Filter:
+    def __init__(self, spec):
+        self.filter = filt.parse(spec)
+
+    def response(self, flow):
+        if flow.match(self.filter):
+            print("Flow matches filter:")
+            print(flow)
 
 
 def start():
     if len(sys.argv) != 2:
         raise ValueError("Usage: -s 'filt.py FILTER'")
-    state["filter"] = filt.parse(sys.argv[1])
-
-
-def response(flow):
-    if flow.match(state["filter"]):
-        print("Flow matches filter:")
-        print(flow)
+    return Filter(sys.argv[1])
diff --git a/examples/flowwriter.py b/examples/flowwriter.py
index 07c7ca20..df2e5a40 100644
--- a/examples/flowwriter.py
+++ b/examples/flowwriter.py
@@ -3,20 +3,21 @@ import sys
 
 from mitmproxy.flow import FlowWriter
 
-state = {}
+
+class Writer:
+    def __init__(self, path):
+        if path == "-":
+            f = sys.stdout
+        else:
+            f = open(path, "wb")
+        self.w = FlowWriter(f)
+
+    def response(self, flow):
+        if random.choice([True, False]):
+            self.w.add(flow)
 
 
 def start():
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "flowriter.py filename"')
-
-    if sys.argv[1] == "-":
-        f = sys.stdout
-    else:
-        f = open(sys.argv[1], "wb")
-    state["flow_writer"] = FlowWriter(f)
-
-
-def response(flow):
-    if random.choice([True, False]):
-        state["flow_writer"].add(flow)
+    return Writer(sys.argv[1])
diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py
index 352c3c24..33d18bbd 100644
--- a/examples/iframe_injector.py
+++ b/examples/iframe_injector.py
@@ -3,26 +3,27 @@
 import sys
 from bs4 import BeautifulSoup
 
-iframe_url = None
+
+class Injector:
+    def __init__(self, iframe_url):
+        self.iframe_url = iframe_url
+
+    def response(self, flow):
+        if flow.request.host in self.iframe_url:
+            return
+        html = BeautifulSoup(flow.response.content, "lxml")
+        if html.body:
+            iframe = html.new_tag(
+                "iframe",
+                src=self.iframe_url,
+                frameborder=0,
+                height=0,
+                width=0)
+            html.body.insert(0, iframe)
+            flow.response.content = str(html).encode("utf8")
 
 
 def start():
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "iframe_injector.py url"')
-    global iframe_url
-    iframe_url = sys.argv[1]
-
-
-def response(flow):
-    if flow.request.host in iframe_url:
-        return
-    html = BeautifulSoup(flow.response.content, "lxml")
-    if html.body:
-        iframe = html.new_tag(
-            "iframe",
-            src=iframe_url,
-            frameborder=0,
-            height=0,
-            width=0)
-        html.body.insert(0, iframe)
-        flow.response.content = str(html).encode("utf8")
+    return Injector(sys.argv[1])
diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index ab068e47..ed5f2ecd 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,18 +87,27 @@ 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):
     def __init__(self, callback):
         self.callback = callback
 
+    def filter(self, event):
+        if event.is_directory:
+            return False
+        if os.path.basename(event.src_path).startswith("."):
+            return False
+        return True
+
     def on_modified(self, event):
-        self.callback()
+        if self.filter(event):
+            self.callback()
 
     def on_created(self, event):
-        self.callback()
+        if self.filter(event):
+            self.callback()
 
 
 class Script:
@@ -118,27 +140,32 @@ 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)
+        ret = self.run("start")
+        if ret:
+            self.ns = ret
+            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..826d2a91 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.start
 
 
 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):
@@ -116,6 +116,19 @@ class TestScript(mastertest.MasterTest):
         assert not fm.state.view[0].request.is_replay
         assert fm.state.view[1].request.is_replay
 
+    def test_addon(self):
+        s = state.State()
+        m = master.FlowMaster(options.Options(), None, s)
+        sc = script.Script(
+            tutils.test_data.path(
+                "data/addonscripts/addon.py"
+            )
+        )
+        m.addons.add(sc)
+        assert sc.ns.event_log == [
+            'scriptstart', 'addonstart', 'addonconfigure'
+        ]
+
 
 class TestScriptLoader(mastertest.MasterTest):
     def test_simple(self):
diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py
new file mode 100644
index 00000000..3b09d231
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/addon.py
@@ -0,0 +1,22 @@
+event_log = []
+
+
+class Addon:
+    @property
+    def event_log(self):
+        return event_log
+
+    def start(self):
+        event_log.append("addonstart")
+
+    def configure(self, options):
+        event_log.append("addonconfigure")
+
+
+def configure(options):
+    event_log.append("addonconfigure")
+
+
+def start():
+    event_log.append("scriptstart")
+    return Addon()
diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py
index b6ac8d89..890e6f4e 100644
--- a/test/mitmproxy/data/addonscripts/recorder.py
+++ b/test/mitmproxy/data/addonscripts/recorder.py
@@ -2,24 +2,24 @@ from mitmproxy import controller
 from mitmproxy import ctx
 import sys
 
-call_log = []
 
-if len(sys.argv) > 1:
-    name = sys.argv[1]
-else:
-    name = "solo"
+class CallLogger:
+    call_log = []
 
-# Keep a log of all possible event calls
-evts = list(controller.Events) + ["configure"]
-for i in evts:
-    def mkprox():
-        evt = i
+    def __init__(self, name = "solo"):
+        self.name = name
 
-        def prox(*args, **kwargs):
-            lg = (name, evt, args, kwargs)
-            if evt != "log":
-                ctx.log.info(str(lg))
-            call_log.append(lg)
-            ctx.log.debug("%s %s" % (name, evt))
-        return prox
-    globals()[i] = mkprox()
+    def __getattr__(self, attr):
+        if attr in controller.Events:
+            def prox(*args, **kwargs):
+                lg = (self.name, attr, args, kwargs)
+                if attr != "log":
+                    ctx.log.info(str(lg))
+                self.call_log.append(lg)
+                ctx.log.debug("%s %s" % (self.name, attr))
+            return prox
+        raise AttributeError
+
+
+def start():
+    return CallLogger(*sys.argv[1:])