From 9450a528fd4fa25b00560ba54a0428bb2191a318 Mon Sep 17 00:00:00 2001 From: "gergely.erdelyi" Date: Sat, 26 Jul 2008 10:18:28 +0000 Subject: [PATCH] init.py: Implemented a simple watchdog to catch runaway scripts --- python/init.py | 62 +++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 61 insertions(+), 1 deletion(-) diff --git a/python/init.py b/python/init.py index e091691..9521a04 100644 --- a/python/init.py +++ b/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"