From 3f6521f9125b13f61192fefde94d9b6d0fc6e489 Mon Sep 17 00:00:00 2001
From: Matthew Shao <me@matshao.com>
Date: Wed, 11 Nov 2015 09:03:05 +0800
Subject: [PATCH] Use watchdog to reload scripts automatically.

---
 libmproxy/console/__init__.py |  1 +
 libmproxy/flow.py             |  4 ++++
 libmproxy/script.py           | 22 ++++++++++++++++++++++
 3 files changed, 27 insertions(+)

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 3bc0c0912..7ef1185ce 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -225,6 +225,7 @@ class ConsoleMaster(flow.FlowMaster):
                 if err:
                     print >> sys.stderr, "Script load error:", err
                     sys.exit(1)
+                script.ObserveScripts(self, i)
 
         if options.outfile:
             err = self.start_stream_to_path(
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 55a4dbcfb..7cce5193e 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -667,6 +667,10 @@ class FlowMaster(controller.Master):
             self.add_event("Script error:\n" + str(e), "error")
         self.scripts.remove(script_obj)
 
+    def reload_scripts(self):
+        for s in self.scripts[:]:
+            s.load()
+    
     def load_script(self, command):
         """
             Loads a script. Returns an error description if something went
diff --git a/libmproxy/script.py b/libmproxy/script.py
index 9d051c129..896f52451 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -4,6 +4,9 @@ import traceback
 import threading
 import shlex
 import sys
+from watchdog.observers import Observer
+from watchdog.events import PatternMatchingEventHandler, FileModifiedEvent
+from .console import signals
 
 
 class ScriptError(Exception):
@@ -192,3 +195,22 @@ def concurrent(fn):
         return _concurrent
     raise NotImplementedError(
         "Concurrent decorator not supported for '%s' method." % fn.func_name)
+
+
+class ScriptModified(PatternMatchingEventHandler):
+    
+    def __init__(self, FlowMaster):
+        self.FlowMaster = FlowMaster
+        PatternMatchingEventHandler.__init__(self, ignore_directories=True, patterns=["*.py"])
+    
+    def on_modified(self, event=FileModifiedEvent):
+        self.FlowMaster.reload_scripts()
+        signals.status_message.send(message="script: <{0}> reloaded.".format(event.src_path))
+
+
+def ObserveScripts(FlowMaster, path):
+    script_dir = os.path.dirname(path)
+    event_handler = ScriptModified(FlowMaster)
+    observer = Observer()
+    observer.schedule(event_handler, script_dir)
+    observer.start()