import sys import ops.cmd import ops.env import ops.db import ops from util.DSZPyLogger import getLogger logger = getLogger('SAFETY') unhookables = ['dir', 'registryquery', 'registryadd', 'get', 'put', 'run', 'dll_u', 'injectdll', 'cd', 'copy', 'move', 'delete'] warnhookables = ['eventlogedit'] def ensureTable(dbHandle=None): if (dbHandle is None): dbHandle = ops.db.Database(db=ops.db.TARGET_DB, isolation_level=None) curs = dbHandle.connection.cursor() else: curs = dbHandle.connection.cursor() try: curs.execute('CREATE TABLE safetyhandlers (plugin, handlerfunc)') except: pass return curs def cmdenv(plugin): return ops.env.get(('OPS_SAFE_%s' % plugin)) def getSafetyHandlerNames(plugin): retval = [] if (cmdenv(plugin) is not None): handlerstr = cmdenv(plugin) handlernames = handlerstr.split(',') retval = [] for handlername in handlernames: lastdot = handlername.rindex('.') modname = handlername[0:lastdot] funcname = handlername[(lastdot + 1):] retval.append((modname, funcname)) return retval def getSafetyHandlerFuncs(plugin): funcnames = getSafetyHandlerNames(plugin) retval = [] for funcname in funcnames: try: stepimport(funcname[0]) retval.append(sys.modules[funcname[0]].__dict__[funcname[1]]) except ImportError as ex: logger.log(10, ('Could not find module %s' % funcname[0])) raise ex except KeyError as ex: logger.log(10, ('Could not find function %s in module %s' % (funcname[1], funcname[0]))) raise ex return retval def addSafetyHandler(plugin, fullfunc): current = getSafetyHandlerNames(plugin) lastdot = fullfunc.rindex('.') modname = fullfunc[0:lastdot] funcname = fullfunc[(lastdot + 1):] splitfunc = (modname, funcname) if (splitfunc not in current): current.append(splitfunc) writeSafetyHandlers(plugin, current) def clearSafetyHandler(plugin): writeSafetyHandlers(plugin, None) def removeSafetyHandler(plugin, fullfunc): current = getSafetyHandlerNames(plugin) lastdot = fullfunc.rindex('.') modname = fullfunc[0:lastdot] funcname = fullfunc[(lastdot + 1):] splitfunc = (modname, funcname) for handler in current: if (handler == splitfunc): current.remove(handler) writeSafetyHandlers(plugin, current) GENERIC_WRAPPER_SCRIPT = 'wrappers/safetywrapper.py' def writeSafetyHandlers(plugin, handlers): if ((handlers is None) or (len(handlers) == 0)): ops.env.delete(('OPS_SAFE_%s' % plugin)) ops.cmd.getDszCommand('wrappers', unregister=plugin, script=GENERIC_WRAPPER_SCRIPT, project='Ops', pre=True).execute() return else: strval = ('%s.%s' % handlers[0]) for handler in handlers[1:]: strval += (',%s.%s' % handler) ops.env.set(('OPS_SAFE_%s' % plugin), strval) if (len(handlers) == 1): ops.cmd.getDszCommand('wrappers', register=plugin, script=GENERIC_WRAPPER_SCRIPT, project='Ops', pre=True).execute() ops.warn(('%d safety %s registered for %s' % (len(handlers), ('handler' if (len(handlers) == 1) else 'handlers'), plugin))) def loadHandlers(): tdb = ops.db.get_tdb() ensureTable(tdb) curs = tdb.connection.execute('SELECT * FROM safetyhandlers') for row in curs: for handler in row['handlerfunc'].split(','): addSafetyHandler(row['plugin'], handler) def saveHandlers(): tdb = ops.db.get_tdb() ensureTable(tdb) tdb.connection.execute('DELETE FROM safetyhandlers') env = ops.cmd.quickrun('lpgetenv') for var in env.envitem: if (var.option[0:9] == 'OPS_SAFE_'): tdb.connection.execute('INSERT INTO safetyhandlers(plugin, handlerfunc) VALUES(?, ?)', (var.option[9:], var.value)) def listSafetyHandlers(): pass def stepimport(modname): packs = modname.split('.') curpack = packs[0] __import__(curpack) for pack in packs[1:]: curpack += ('.' + pack) __import__(curpack) def doSafetyHandlers(commandobj): if (not isinstance(commandobj, ops.cmd.DszCommand)): raise RuntimeError, ('Script did not provide an ops.cmd.DszCommand instance; instead it gave me a %s' % type(commandobj)) good = True msgparts = [] if (not commandobj.validateInput()): good = False msgparts.append('Input validation failed') handlers = getSafetyHandlerFuncs(commandobj.plugin) for handler in handlers: try: (isgood, msg) = handler(commandobj) except AttributeError: (isgood, msg) = (True, '') if (not isgood): good = False if (msg != ''): msgparts.append(msg) return (good, msgparts) def _usage(): return 'Usage: \n\tsafetychecks.py \n\tsafetychecks.py clear \n\tsafetychecks.py \n\n\tNote: name of safety check is not required for clear' if ((__name__ == '__main__') or (__name__ == '__ops_survey_plugin__')): good = False if (len(sys.argv) == 2): action = sys.argv[1] if (action == 'load'): loadHandlers() ops.info('Loaded safety handlers from previous op(s)') good = True elif (action == 'save'): saveHandlers() ops.info('Saved safety handlers for future op(s)') good = True elif (action == 'list'): listSafetyHandlers() good = True elif (len(sys.argv) == 3): action = sys.argv[1] plugin = sys.argv[2] if (action == 'clear'): clearSafetyHandler(plugin) good = True elif (len(sys.argv) == 4): action = sys.argv[1] plugin = sys.argv[2] handler = sys.argv[3] if (action == 'add'): addSafetyHandler(plugin, handler) good = True elif (action == 'delete'): removeSafetyHandler(plugin, handler) good = True if (not good): ops.warn('You called this wrong! Are you sure you should be doing this?') ops.info(_usage())