diff options
| author | Maximilian Hils <git@maximilianhils.com> | 2015-11-14 04:21:38 +0100 |
|---|---|---|
| committer | Maximilian Hils <git@maximilianhils.com> | 2015-11-14 04:21:38 +0100 |
| commit | 4499ab61c04d765fc191227e25af0ee1cc98a83a (patch) | |
| tree | 3208cc7176104c7ee9fcd7173719ab744b486f41 /libmproxy/script/script.py | |
| parent | 247f27d8219bb63ef64dd33935e222fdad507631 (diff) | |
| parent | cd0b9e01be4041275165ab5e90b524ab0b3247f0 (diff) | |
| download | mitmproxy-4499ab61c04d765fc191227e25af0ee1cc98a83a.tar.gz mitmproxy-4499ab61c04d765fc191227e25af0ee1cc98a83a.tar.bz2 mitmproxy-4499ab61c04d765fc191227e25af0ee1cc98a83a.zip | |
structure libmproxy.script
Diffstat (limited to 'libmproxy/script/script.py')
| -rw-r--r-- | libmproxy/script/script.py | 123 |
1 files changed, 123 insertions, 0 deletions
diff --git a/libmproxy/script/script.py b/libmproxy/script/script.py new file mode 100644 index 00000000..a58ba0af --- /dev/null +++ b/libmproxy/script/script.py @@ -0,0 +1,123 @@ +""" +The script object representing mitmproxy inline scripts. +Script objects know nothing about mitmproxy or mitmproxy's API - this knowledge is provided +by the mitmproxy-specific ScriptContext. +""" +from __future__ import absolute_import, print_function, division +import os +import shlex +import traceback +import sys +import blinker + +from watchdog.events import PatternMatchingEventHandler, FileModifiedEvent +from watchdog.observers import Observer + +from ..exceptions import ScriptException + +script_change = blinker.Signal() + + +class Script(object): + """ + Script object representing an inline script. + """ + + def __init__(self, command, context, use_reloader=True): + self.command = command + self.args = self.parse_command(command) + self.ctx = context + self.ns = None + self.load() + if use_reloader: + self.start_observe() + + @classmethod + def parse_command(cls, command): + if not command or not command.strip(): + raise ScriptException("Empty script command.") + if os.name == "nt": # Windows: escape all backslashes in the path. + backslashes = shlex.split(command, posix=False)[0].count("\\") + command = command.replace("\\", "\\\\", backslashes) + args = shlex.split(command) + args[0] = os.path.expanduser(args[0]) + if not os.path.exists(args[0]): + raise ScriptException( + ("Script file not found: %s.\r\n" + "If your script path contains spaces, " + "make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") % + args[0]) + elif os.path.isdir(args[0]): + raise ScriptException("Not a file: %s" % args[0]) + return args + + def load(self): + """ + Loads an inline script. + + Returns: + The return value of self.run("start", ...) + + Raises: + ScriptException on failure + """ + if self.ns is not None: + self.unload() + script_dir = os.path.dirname(os.path.abspath(self.args[0])) + ns = {'__file__': os.path.abspath(self.args[0])} + sys.path.append(script_dir) + try: + execfile(self.args[0], ns, ns) + except Exception as e: + # Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/ + raise ScriptException(traceback.format_exc(e)) + sys.path.pop() + self.ns = ns + return self.run("start", self.args) + + def unload(self): + ret = self.run("done") + self.ns = None + return ret + + def run(self, name, *args, **kwargs): + """ + Runs an inline script hook. + + Returns: + The return value of the method. + None, if the script does not provide the method. + + Raises: + ScriptException if there was an exception. + """ + f = self.ns.get(name) + if f: + try: + return f(self.ctx, *args, **kwargs) + except Exception as e: + raise ScriptException(traceback.format_exc(e)) + else: + return None + + def start_observe(self): + script_dir = os.path.dirname(self.args[0]) + event_handler = ScriptModified(self) + observer = Observer() + observer.schedule(event_handler, script_dir) + observer.start() + + def stop_observe(self): + raise NotImplementedError() # FIXME + + +class ScriptModified(PatternMatchingEventHandler): + def __init__(self, script): + # We could enumerate all relevant *.py files (as werkzeug does it), + # but our case looks like it isn't as simple as enumerating sys.modules. + # This should be good enough for now. + super(ScriptModified, self).__init__(ignore_directories=True, patterns=["*.py"]) + self.script = script + + def on_modified(self, event=FileModifiedEvent): + script_change.send(self.script) |