Move script.Context to flow.ScriptContext

3 files changed, 27 insertions, 18 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 25fa3fc1..e649fbb8 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -15,6 +15,21 @@ class RunException(Exception):
         self.errout = errout
 
 
+class ScriptContext:
+    def __init__(self, master):
+        self._master = master
+
+    def log(self, *args, **kwargs):
+        """
+            Logs an event. 
+            
+            How this is handled depends on the front-end. mitmdump will display
+            events if the eventlog flag ("-e") was passed. mitmproxy sends
+            output to the eventlog for display ("v" keyboard shortcut).
+        """
+        self._master.add_event(*args, **kwargs)
+
+
 class Headers:
     def __init__(self, lst=None):
         if lst:
@@ -1055,7 +1070,6 @@ class State(object):
             i.kill(master)
 
 
-
 class FlowMaster(controller.Master):
     def __init__(self, server, state):
         controller.Master.__init__(self, server)
@@ -1086,7 +1100,7 @@ class FlowMaster(controller.Master):
         """
             Returns an (error, script) tuple.
         """
-        s = script.Script(path, self)
+        s = script.Script(path, ScriptContext(self))
         try:
             s.load()
         except script.ScriptError, v:
diff --git a/libmproxy/script.py b/libmproxy/script.py
index 03eff958..251d6dea 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -1,15 +1,8 @@
-import imp, os, traceback
+import imp, os, traceback, flow
 
 class ScriptError(Exception):
     pass
 
-class Context:
-    def __init__(self, master, state):
-        self.master, self.state = master, state
-
-    def log(self, *args, **kwargs):
-        self.master.add_event(*args, **kwargs)
-
 
 class Script:
     """
@@ -19,9 +12,8 @@ class Script:
             s.load()
             s.run("start")
     """
-    def __init__(self, path, master):
-        self.path = path
-        self.ctx = Context(master, master.state)
+    def __init__(self, path, ctx):
+        self.path, self.ctx = path, ctx
         self.mod = None
         self.ns = None
 
diff --git a/test/test_script.py b/test/test_script.py
index 9a1c3062..32ee2d95 100644
--- a/test/test_script.py
+++ b/test/test_script.py
@@ -6,8 +6,9 @@ class uScript(libpry.AutoTree):
     def test_simple(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
+        ctx = flow.ScriptContext(fm)
 
-        p = script.Script(os.path.join("scripts", "a.py"), fm)
+        p = script.Script(os.path.join("scripts", "a.py"), ctx)
         p.load()
         assert "here" in p.ns
         assert p.run("here") == (True, 1)
@@ -24,26 +25,28 @@ class uScript(libpry.AutoTree):
     def test_err(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
+        ctx = flow.ScriptContext(fm)
 
-        s = script.Script("nonexistent", fm)
+
+        s = script.Script("nonexistent", ctx)
         libpry.raises(
             "no such file",
             s.load
         )
 
-        s = script.Script("scripts", fm)
+        s = script.Script("scripts", ctx)
         libpry.raises(
             "not a file",
             s.load
         )
 
-        s = script.Script("scripts/syntaxerr.py", fm)
+        s = script.Script("scripts/syntaxerr.py", ctx)
         libpry.raises(
             script.ScriptError,
             s.load
         )
 
-        s = script.Script("scripts/loaderr.py", fm)
+        s = script.Script("scripts/loaderr.py", ctx)
         libpry.raises(
             script.ScriptError,
             s.load