diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/flow.py | 25 | ||||
| -rw-r--r-- | libmproxy/proxy/server.py | 12 | ||||
| -rw-r--r-- | libmproxy/script.py | 61 |
3 files changed, 54 insertions, 44 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 16340d5d..cbb22c3a 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/proxy/server.py b/libmproxy/proxy/server.py index 2711bd0e..e77439fb 100644 --- a/libmproxy/proxy/server.py +++ b/libmproxy/proxy/server.py @@ -167,12 +167,12 @@ class ConnectionHandler: self.channel.tell("serverdisconnect", self) self.server_conn = None - def set_server_address(self, address): + def set_server_address(self, addr): """ Sets a new server address with the given priority. Does not re-establish either connection or SSL handshake. """ - address = tcp.Address.wrap(address) + address = tcp.Address.wrap(addr) # Don't reconnect to the same destination. if self.server_conn and self.server_conn.address == address: @@ -309,15 +309,15 @@ class ConnectionHandler: self.client_conn.finish() def log(self, msg, level, subs=()): - msg = [ + full_msg = [ "%s:%s: %s" % (self.client_conn.address.host, self.client_conn.address.port, msg)] for i in subs: - msg.append(" -> " + i) - msg = "\n".join(msg) - self.channel.tell("log", Log(msg, level)) + full_msg.append(" -> " + i) + full_msg = "\n".join(full_msg) + self.channel.tell("log", Log(full_msg, level)) def find_cert(self): host = self.server_conn.address.host 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) |