diff options
| -rw-r--r-- | libmproxy/flow.py | 25 | ||||
| -rw-r--r-- | libmproxy/script.py | 61 | ||||
| -rw-r--r-- | test/scripts/a.py | 5 | ||||
| -rw-r--r-- | test/scripts/a_helper.py | 4 | ||||
| -rw-r--r-- | test/scripts/unloaderr.py | 2 | ||||
| -rw-r--r-- | test/test_examples.py | 2 | ||||
| -rw-r--r-- | test/test_script.py | 224 |
7 files changed, 170 insertions, 153 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 1a052f51..1bc73d91 100644 --- a/libmproxy/flow.py +++ b/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 diff --git a/libmproxy/script.py b/libmproxy/script.py index 6167ae05..e13f0e2b 100644 --- a/libmproxy/script.py +++ b/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. + + Returns: + The return value of self.run("start", ...) - Raises ScriptError on failure, with argument equal to an error - message that may be a formatted traceback. + 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) diff --git a/test/scripts/a.py b/test/scripts/a.py index 210fea78..d4272ac8 100644 --- a/test/scripts/a.py +++ b/test/scripts/a.py @@ -1,7 +1,4 @@ -import argparse - -parser = argparse.ArgumentParser() -parser.add_argument('--var', type=int) +from a_helper import parser var = 0 diff --git a/test/scripts/a_helper.py b/test/scripts/a_helper.py new file mode 100644 index 00000000..2eeed0d4 --- /dev/null +++ b/test/scripts/a_helper.py @@ -0,0 +1,4 @@ +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument('--var', type=int)
\ No newline at end of file diff --git a/test/scripts/unloaderr.py b/test/scripts/unloaderr.py new file mode 100644 index 00000000..f3743107 --- /dev/null +++ b/test/scripts/unloaderr.py @@ -0,0 +1,2 @@ +def done(ctx): + raise RuntimeError()
\ No newline at end of file diff --git a/test/test_examples.py b/test/test_examples.py index e9bccd2e..54577e67 100644 --- a/test/test_examples.py +++ b/test/test_examples.py @@ -22,7 +22,7 @@ def test_load_scripts(): try: s = script.Script(f, tmaster) # Loads the script file. except Exception as v: - if not "ImportError" in str(v): + if "ImportError" not in str(v): raise else: s.unload() diff --git a/test/test_script.py b/test/test_script.py index 0a063740..1b0e5a5b 100644 --- a/test/test_script.py +++ b/test/test_script.py @@ -1,120 +1,124 @@ -from libmproxy import script, flow -import tutils -import shlex import os import time import mock +from libmproxy import script, flow +import tutils + + +def test_simple(): + s = flow.State() + fm = flow.FlowMaster(None, s) + sp = tutils.test_data.path("scripts/a.py") + p = script.Script("%s --var 40" % sp, fm) + + assert "here" in p.ns + assert p.run("here") == 41 + assert p.run("here") == 42 + + tutils.raises(script.ScriptError, p.run, "errargs") + + # Check reload + p.load() + assert p.run("here") == 41 + + +def test_duplicate_flow(): + s = flow.State() + fm = flow.FlowMaster(None, s) + fm.load_script(tutils.test_data.path("scripts/duplicate_flow.py")) + f = tutils.tflow() + fm.handle_request(f) + assert fm.state.flow_count() == 2 + assert not fm.state.view[0].request.is_replay + assert fm.state.view[1].request.is_replay + + +def test_err(): + s = flow.State() + fm = flow.FlowMaster(None, s) + + tutils.raises( + "not found", + script.Script, "nonexistent", fm + ) + + tutils.raises( + "not a file", + script.Script, tutils.test_data.path("scripts"), fm + ) + + tutils.raises( + script.ScriptError, + script.Script, tutils.test_data.path("scripts/syntaxerr.py"), fm + ) + + tutils.raises( + script.ScriptError, + script.Script, tutils.test_data.path("scripts/loaderr.py"), fm + ) + + scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), fm) + tutils.raises(script.ScriptError, scr.unload) -class TestScript: - def test_simple(self): - s = flow.State() - fm = flow.FlowMaster(None, s) - sp = tutils.test_data.path("scripts/a.py") - p = script.Script("%s --var 40" % sp, fm) - - assert "here" in p.ns - assert p.run("here") == (True, 41) - assert p.run("here") == (True, 42) - - ret = p.run("errargs") - assert not ret[0] - assert len(ret[1]) == 2 - - # Check reload - p.load() - assert p.run("here") == (True, 41) - - def test_duplicate_flow(self): - s = flow.State() - fm = flow.FlowMaster(None, s) - fm.load_script(tutils.test_data.path("scripts/duplicate_flow.py")) - f = tutils.tflow() - fm.handle_request(f) - assert fm.state.flow_count() == 2 - assert not fm.state.view[0].request.is_replay - assert fm.state.view[1].request.is_replay - - def test_err(self): - s = flow.State() - fm = flow.FlowMaster(None, s) - - tutils.raises( - "not found", - script.Script, "nonexistent", fm - ) - - tutils.raises( - "not a file", - script.Script, tutils.test_data.path("scripts"), fm - ) - - tutils.raises( - script.ScriptError, - script.Script, tutils.test_data.path("scripts/syntaxerr.py"), fm - ) - - tutils.raises( - script.ScriptError, - script.Script, tutils.test_data.path("scripts/loaderr.py"), fm - ) - - def test_concurrent(self): - s = flow.State() - fm = flow.FlowMaster(None, s) - fm.load_script(tutils.test_data.path("scripts/concurrent_decorator.py")) - - with mock.patch("libmproxy.controller.DummyReply.__call__") as m: - f1, f2 = tutils.tflow(), tutils.tflow() - t_start = time.time() - fm.handle_request(f1) - f1.reply() - fm.handle_request(f2) - f2.reply() - - # Two instantiations - assert m.call_count == 0 # No calls yet. - assert (time.time() - t_start) < 0.09 - - def test_concurrent2(self): - s = flow.State() - fm = flow.FlowMaster(None, s) - s = script.Script( - tutils.test_data.path("scripts/concurrent_decorator.py"), - fm) - s.load() - m = mock.Mock() - - class Dummy: - def __init__(self): - self.response = self - self.error = self - self.reply = m +def test_concurrent(): + s = flow.State() + fm = flow.FlowMaster(None, s) + fm.load_script(tutils.test_data.path("scripts/concurrent_decorator.py")) + with mock.patch("libmproxy.controller.DummyReply.__call__") as m: + f1, f2 = tutils.tflow(), tutils.tflow() t_start = time.time() + fm.handle_request(f1) + f1.reply() + fm.handle_request(f2) + f2.reply() + + # Two instantiations + assert m.call_count == 0 # No calls yet. + assert (time.time() - t_start) < 0.09 + - for hook in ("clientconnect", - "serverconnect", - "response", - "error", - "clientconnect"): - d = Dummy() - assert s.run(hook, d)[0] - d.reply() - while (time.time() - t_start) < 20 and m.call_count <= 5: - if m.call_count == 5: - return - time.sleep(0.001) - assert False - - def test_concurrent_err(self): - s = flow.State() - fm = flow.FlowMaster(None, s) - tutils.raises( - "decorator not supported for this method", - script.Script, - tutils.test_data.path("scripts/concurrent_decorator_err.py"), - fm) +def test_concurrent2(): + s = flow.State() + fm = flow.FlowMaster(None, s) + s = script.Script( + tutils.test_data.path("scripts/concurrent_decorator.py"), + fm) + s.load() + m = mock.Mock() + + class Dummy: + def __init__(self): + self.response = self + self.error = self + self.reply = m + + t_start = time.time() + + for hook in ("clientconnect", + "serverconnect", + "response", + "error", + "clientconnect"): + d = Dummy() + s.run(hook, d) + d.reply() + while (time.time() - t_start) < 20 and m.call_count <= 5: + if m.call_count == 5: + return + time.sleep(0.001) + assert False + + +def test_concurrent_err(): + s = flow.State() + fm = flow.FlowMaster(None, s) + tutils.raises( + "Concurrent decorator not supported for 'start' method", + script.Script, + tutils.test_data.path("scripts/concurrent_decorator_err.py"), + fm) def test_command_parsing(): @@ -122,4 +126,4 @@ def test_command_parsing(): fm = flow.FlowMaster(None, s) absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py")) s = script.Script(absfilepath, fm) - assert os.path.isfile(s.argv[0]) + assert os.path.isfile(s.args[0]) |