177 lines
No EOL
6.3 KiB
Python
177 lines
No EOL
6.3 KiB
Python
|
|
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 <load|save|list>\n\tsafetychecks.py clear <name of plugin>\n\tsafetychecks.py <add|delete> <name of safety check>\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()) |