Start refactoring scripts

- Move ScriptContext into script module
- Use mock module instead of hand-rolled mock objects in tests

4 files changed, 45 insertions, 42 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 210e2b95..b8d4a105 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -781,6 +781,9 @@ class ConsoleMaster(flow.FlowMaster):
         else:
             self.view_flowlist()
 
+    def edit_scripts(self, *args, **kwargs):
+        pass
+
     def loop(self):
         changed = True
         try:
@@ -883,7 +886,7 @@ class ConsoleMaster(flow.FlowMaster):
                                     grideditor.ScriptEditor(
                                         self,
                                         [[i.argv[0]] for i in self.scripts],
-                                        None
+                                        self.edit_scripts
                                     )
                                 )
                                 #if self.scripts:
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 32306513..e88bf985 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -143,39 +143,6 @@ class SetHeaders:
                     f.request.headers.add(header, value)
 
 
-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)
-
-    def duplicate_flow(self, f):
-        """
-            Returns a duplicate of the specified flow. The flow is also
-            injected into the current state, and is ready for editing, replay,
-            etc.
-        """
-        self._master.pause_scripts = True
-        f = self._master.duplicate_flow(f)
-        self._master.pause_scripts = False
-        return f
-
-    def replay_request(self, f):
-        """
-            Replay the request on the current flow. The response will be added
-            to the flow object.
-        """
-        self._master.replay_request(f)
-
-
 class decoded(object):
     """
 
@@ -1431,14 +1398,14 @@ class FlowMaster(controller.Master):
         """
             Returns an (error, script) tuple.
         """
-        s = script.Script(script_argv, ScriptContext(self))
+        s = script.Script(script_argv, self)
         try:
             s.load()
         except script.ScriptError, v:
             return (v.args[0], None)
         return (None, s)
 
-    def unload_script(self,script):
+    def unload_script(self, script):
         script.unload()
         self.scripts.remove(script)
 
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 15f75b8d..1894f7f0 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -110,15 +110,15 @@ class HandleSNI:
         self.handler, self.client_conn, self.host, self.port = handler, client_conn, host, port
         self.cert, self.key = cert, key
 
-    def __call__(self, connection):
+    def __call__(self, client_connection):
         try:
-            sn = connection.get_servername()
+            sn = client_connection.get_servername()
             if sn:
                 self.handler.get_server_connection(self.client_conn, "https", self.host, self.port, sn)
                 new_context = SSL.Context(SSL.TLSv1_METHOD)
                 new_context.use_privatekey_file(self.key)
                 new_context.use_certificate(self.cert.x509)
-                connection.set_context(new_context)
+                client_connection.set_context(new_context)
                 self.handler.sni = sn.decode("utf8").encode("idna")
         # An unhandled exception in this method will core dump PyOpenSSL, so
         # make dang sure it doesn't happen.
diff --git a/libmproxy/script.py b/libmproxy/script.py
index f8a0d085..0f1056f6 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -5,6 +5,39 @@ class ScriptError(Exception):
     pass
 
 
+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)
+
+    def duplicate_flow(self, f):
+        """
+            Returns a duplicate of the specified flow. The flow is also
+            injected into the current state, and is ready for editing, replay,
+            etc.
+        """
+        self._master.pause_scripts = True
+        f = self._master.duplicate_flow(f)
+        self._master.pause_scripts = False
+        return f
+
+    def replay_request(self, f):
+        """
+            Replay the request on the current flow. The response will be added
+            to the flow object.
+        """
+        self._master.replay_request(f)
+
+
 class Script:
     """
         The instantiator should do something along this vein:
@@ -12,9 +45,9 @@ class Script:
             s = Script(argv, master)
             s.load()
     """
-    def __init__(self, argv, ctx):
+    def __init__(self, argv, master):
         self.argv = argv
-        self.ctx = ctx
+        self.ctx = ScriptContext(master)
         self.ns = None
 
     def load(self):
@@ -82,4 +115,4 @@ def concurrent(fn):
         def _concurrent(ctx, conn):
             _handle_concurrent_reply(fn, conn, [ctx, conn])
         return _concurrent
-    raise NotImplementedError("Concurrent decorator not supported for this method.")
\ No newline at end of file
+    raise NotImplementedError("Concurrent decorator not supported for this method.")