init.py: Implemented a simple watchdog to catch runaway scripts

python/init.py

@@ -11,7 +11,12 @@
 #------------------------------------------------------------
 # init.py - Essential init routines
 #------------------------------------------------------------
-import sys, os, os.path, traceback, warnings
+import os
+import sys
+import time
+import traceback
+import warnings
+
 import _idaapi
 
 # FIXME: Should fix the offending constant instead
@@ -126,6 +131,61 @@ class ScriptBox(Choose):
 
 ScriptBox_instance = ScriptBox([])
 
+#-------------------------------------------------------------
+# Watchdog to catch runaway scripts after a specified timeout
+#
+# Usage:
+#        watchdog.install()
+#        watchdog.activate(10) # Use 10-second timeout
+#
+# Note: The watchdog only works for code running inside
+#       functions, not in global/module namespace.         
+#-------------------------------------------------------------
+class WatchDog():
+    """
+    Python tracer-based watchdog class
+    """
+    def __init__(self, timeout=10):
+        self.timeout = timeout
+        self.installed = False
+        self.active = False
+
+    def install(self):
+        """ Install the tracer function, required for the watchdog """
+        if not self.installed:
+            sys.settrace(self.tracer)
+            self.installed = True
+
+    def activate(self, timeout=None):
+        """ Activate the watchdog, with optional timeout change """
+        assert self.installed, "WatchDog must be installed before activating"
+        if timeout:
+            self.timeout = timeout
+        self.reset()
+        self.active = True
+
+    def deactivate(self):
+        """ Deactivate the watchdog """
+        self.active = True
+
+    def reset(self):
+        """ Reset the timer, useful for long-running scripts """
+        self.timestamp = time.clock()
+
+    def tracer(self, frame, event, arg):
+        """ Tracer function that receives the tracing events """
+        if not self.active:
+            return None
+        if event == 'line':
+            if time.clock() - self.timestamp > self.timeout:
+                if AskYN(0, "The script has not finished in %d seconds\nWould you like to stop it now?" % self.timeout) == 1:
+                    raise KeyboardInterrupt
+                else:
+                    self.timestamp = time.clock()
+        return self.tracer
+
+watchdog = WatchDog(10)
+
 # Load the users personal init file
 userrc = get_user_idadir() + os.sep + "idapythonrc.py"