diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py index 6f09a6b32..cef2013e0 100644 --- a/libmproxy/console/__init__.py +++ b/libmproxy/console/__init__.py @@ -732,6 +732,8 @@ class ConsoleMaster(flow.FlowMaster): self.process_flow(f) return f - def script_change(self, script): - self.masterq.put(("script_change", script)) - signals.status_message.send(message="<{}> reloaded.".format(script.args[0])) + def handle_script_change(self, script): + if super(ConsoleMaster, self).handle_script_change(script): + signals.status_message.send(message='"{}" reloaded.'.format(script.filename)) + else: + signals.status_message.send(message='Error reloading "{}".'.format(script.filename)) \ No newline at end of file diff --git a/libmproxy/flow.py b/libmproxy/flow.py index f090d9c6b..97d72992f 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -640,7 +640,6 @@ class FlowMaster(controller.Master): self.stream = None self.apps = AppRegistry() - script.script_change.connect(self.script_change) def start_app(self, host, port): self.apps.add( @@ -664,15 +663,18 @@ class FlowMaster(controller.Master): script_obj.unload() except script.ScriptException as e: self.add_event("Script error:\n" + str(e), "error") + script.reloader.unwatch(script_obj) self.scripts.remove(script_obj) - def load_script(self, command): + def load_script(self, command, use_reloader=True): """ Loads a script. Returns an error description if something went wrong. """ try: s = script.Script(command, script.ScriptContext(self)) + if use_reloader: + script.reloader.watch(s, lambda: self.masterq.put(("script_change", s))) except script.ScriptException as v: return v.args[0] self.scripts.append(s) @@ -1020,8 +1022,33 @@ class FlowMaster(controller.Master): def handle_accept_intercept(self, f): self.state.update_flow(f) - def handle_script_change(self, script): - script.load() + def handle_script_change(self, s): + """ + Handle a script whose contents have been changed on the file system. + + Args: + s (script.Script): the changed script + + Returns: + True, if reloading was successful. + False, otherwise. + """ + ok = True + # We deliberately do not want to fail here. + # In the worst case, we have an "empty" script object. + try: + s.unload() + except script.ScriptException as e: + ok = False + self.add_event('Error reloading "{}": {}'.format(s.filename, str(e))) + try: + s.load() + except script.ScriptException as e: + ok = False + self.add_event('Error reloading "{}": {}'.format(s.filename, str(e))) + else: + self.add_event('"{}" reloaded.'.format(s.filename)) + return ok def shutdown(self): self.unload_scripts() @@ -1039,11 +1066,6 @@ class FlowMaster(controller.Master): self.stream.fo.close() self.stream = None - def script_change(self, script): - self.masterq.put(("script_change", script)) - self.add_event("<{}> reloaded.".format(script.args[0])) - - def read_flows_from_paths(paths): """ Given a list of filepaths, read all flows and return a list of them. diff --git a/libmproxy/script/__init__.py b/libmproxy/script/__init__.py index 0f487795e..8bcdc5a21 100644 --- a/libmproxy/script/__init__.py +++ b/libmproxy/script/__init__.py @@ -1,11 +1,13 @@ -from .script import Script, script_change +from .script import Script from .script_context import ScriptContext from .concurrent import concurrent from ..exceptions import ScriptException +from . import reloader __all__ = [ - "Script", "script_change", + "Script", "ScriptContext", "concurrent", - "ScriptException" + "ScriptException", + "reloader" ] \ No newline at end of file diff --git a/libmproxy/script/reloader.py b/libmproxy/script/reloader.py new file mode 100644 index 000000000..b867238f5 --- /dev/null +++ b/libmproxy/script/reloader.py @@ -0,0 +1,37 @@ +import os +from watchdog.events import PatternMatchingEventHandler +from watchdog.observers import Observer + +_observers = {} + + +def watch(script, callback): + script_dir = os.path.dirname(os.path.abspath(script.args[0])) + event_handler = _ScriptModificationHandler(callback) + observer = Observer() + observer.schedule(event_handler, script_dir) + observer.start() + _observers[script] = observer + + +def unwatch(script): + observer = _observers.pop(script, None) + if observer: + observer.stop() + + +class _ScriptModificationHandler(PatternMatchingEventHandler): + def __init__(self, callback): + # 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(_ScriptModificationHandler, self).__init__( + ignore_directories=True, + patterns=["*.py"] + ) + self.callback = callback + + def on_modified(self, event): + self.callback() + +__all__ = ["watch", "unwatch"] \ No newline at end of file diff --git a/libmproxy/script/script.py b/libmproxy/script/script.py index a58ba0af5..498caf946 100644 --- a/libmproxy/script/script.py +++ b/libmproxy/script/script.py @@ -8,32 +8,27 @@ 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. + Script object representing an inline script. """ - def __init__(self, command, context, use_reloader=True): + def __init__(self, command, context): 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): + @property + def filename(self): + return self.args[0] + + @staticmethod + def parse_command(command): if not command or not command.strip(): raise ScriptException("Empty script command.") if os.name == "nt": # Windows: escape all backslashes in the path. @@ -64,21 +59,22 @@ class Script(object): 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])} + self.ns = {'__file__': os.path.abspath(self.args[0])} sys.path.append(script_dir) try: - execfile(self.args[0], ns, ns) + execfile(self.args[0], self.ns, self.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 + finally: + sys.path.pop() return self.run("start", self.args) def unload(self): - ret = self.run("done") - self.ns = None - return ret + try: + return self.run("done") + finally: + self.ns = None def run(self, name, *args, **kwargs): """ @@ -98,26 +94,4 @@ class Script(object): 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) + return None \ No newline at end of file diff --git a/setup.py b/setup.py index e0a29eefa..f20e26693 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,7 @@ deps = { "six~=1.10.0", "lxml~=3.4.4", "Pillow~=3.0.0", + "watchdog~=0.8.3", } # A script -> additional dependencies dict. scripts = {