improve inline scripts, fix #683, fix #684

libmproxy/flow.py

@@ -659,9 +659,12 @@ class FlowMaster(controller.Master):
         for s in self.scripts[:]:
             self.unload_script(s)
 
-    def unload_script(self, script):
-        script.unload()
-        self.scripts.remove(script)
+    def unload_script(self, script_obj):
+        try:
+            script_obj.unload()
+        except script.ScriptError as e:
+            self.add_event("Script error:\n" + str(e), "error")
+        self.scripts.remove(script_obj)
 
     def load_script(self, command):
         """
@@ -674,16 +677,16 @@ class FlowMaster(controller.Master):
             return v.args[0]
         self.scripts.append(s)
 
-    def run_single_script_hook(self, script, name, *args, **kwargs):
-        if script and not self.pause_scripts:
-            ret = script.run(name, *args, **kwargs)
-            if not ret[0] and ret[1]:
-                e = "Script error:\n" + ret[1][1]
-                self.add_event(e, "error")
+    def _run_single_script_hook(self, script_obj, name, *args, **kwargs):
+        if script_obj and not self.pause_scripts:
+            try:
+                script_obj.run(name, *args, **kwargs)
+            except script.ScriptError as e:
+                self.add_event("Script error:\n" + str(e), "error")
 
     def run_script_hook(self, name, *args, **kwargs):
-        for script in self.scripts:
-            self.run_single_script_hook(script, name, *args, **kwargs)
+        for script_obj in self.scripts:
+            self._run_single_script_hook(script_obj, name, *args, **kwargs)
 
     def get_ignore_filter(self):
         return self.server.config.check_ignore.patterns

libmproxy/script.py

@@ -3,7 +3,7 @@ import os
 import traceback
 import threading
 import shlex
-from . import controller
+import sys
 
 
 class ScriptError(Exception):
@@ -55,21 +55,17 @@ class ScriptContext:
 
 class Script:
     """
-        The instantiator should do something along this vein:
-
-            s = Script(argv, master)
-            s.load()
+        Script object representing an inline script.
     """
 
     def __init__(self, command, master):
-        self.command = command
-        self.argv = self.parse_command(command)
+        self.args = self.parse_command(command)
         self.ctx = ScriptContext(master)
         self.ns = None
         self.load()
 
     @classmethod
-    def parse_command(klass, command):
+    def parse_command(cls, command):
         if not command or not command.strip():
             raise ScriptError("Empty script command.")
         if os.name == "nt":  # Windows: escape all backslashes in the path.
@@ -89,42 +85,52 @@ class Script:
 
     def load(self):
         """
-            Loads a module.
+            Loads an inline script.
 
-            Raises ScriptError on failure, with argument equal to an error
-            message that may be a formatted traceback.
+            Returns:
+                The return value of self.run("start", ...)
+
+            Raises:
+                ScriptError on failure
         """
+        if self.ns is not None:
+            self.unload()
         ns = {}
+        script_dir = os.path.dirname(os.path.abspath(self.args[0]))
+        sys.path.append(script_dir)
         try:
-            execfile(self.argv[0], ns, ns)
-        except Exception as v:
-            raise ScriptError(traceback.format_exc(v))
+            execfile(self.args[0], ns, ns)
+        except Exception as e:
+            # Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/
+            raise ScriptError(traceback.format_exc(e))
+        sys.path.pop()
         self.ns = ns
-        r = self.run("start", self.argv)
-        if not r[0] and r[1]:
-            raise ScriptError(r[1][1])
+        return self.run("start", self.args)
 
     def unload(self):
-        return self.run("done")
+        ret = self.run("done")
+        self.ns = None
+        return ret
 
     def run(self, name, *args, **kwargs):
         """
-            Runs a plugin method.
+            Runs an inline script hook.
 
             Returns:
+                The return value of the method.
+                None, if the script does not provide the method.
 
-                (True, retval) on success.
-                (False, None) on nonexistent method.
-                (False, (exc, traceback string)) if there was an exception.
+            Raises:
+                ScriptError if there was an exception.
         """
         f = self.ns.get(name)
         if f:
             try:
-                return (True, f(self.ctx, *args, **kwargs))
-            except Exception as v:
-                return (False, (v, traceback.format_exc(v)))
+                return f(self.ctx, *args, **kwargs)
+            except Exception as e:
+                raise ScriptError(traceback.format_exc(e))
         else:
-            return (False, None)
+            return None
 
 
 class ReplyProxy(object):
@@ -176,6 +182,7 @@ def concurrent(fn):
             "clientdisconnect"):
         def _concurrent(ctx, obj):
             _handle_concurrent_reply(fn, obj, ctx, obj)
+
         return _concurrent
     raise NotImplementedError(
-        "Concurrent decorator not supported for this method.")
+        "Concurrent decorator not supported for '%s' method." % fn.func_name)