#! /usr/bin/env python3 import sys sys.path.append('modules') sys.path.append('vars') sys.path.append('logger') import time import locale from dialog import Dialog import glob import os import subprocess as sp import tomb import traceback import logger from consolelogger import ConsoleLogger from filelogger import FileLogger from hashlib import sha256 def getModules(): choiceList = [] modules = [] for filename in glob.glob('./modules/*.py'): modules.append(os.path.splitext(os.path.basename(filename))[0]) modules.remove('__init__') for modname in modules: try: module = __import__(modname) instance = module.getInstance() choiceList.append((instance.name,instance.description,False)) del instance del module sys.modules.pop(modname) except Exception as e: print("Exception raised while importing " + modname) return choiceList def prepareModule(moduleName): try: module = __import__(moduleName).getInstance() module.getVars() return module except Exception as e: raise def sha256sum(file): f = open(file, 'rb') result = sha256(f.read()).hexdigest() f.close() return result def bagAndTag(): tombPath = tomb.getPath() hashes = open(tombPath + "hashes.txt",'a+') boneList = [] for root, dirnames, filenames in os.walk(tombPath): for dir in dirnames: for subroot, subdirnames, subfilenames in os.walk(tombPath + dir): for filename in subfilenames: boneList.append(subroot + "/" + filename) d.gauge_start(title="Hashing all collected artifacts...",width=60,height=10) boneCount = 0 for bone in boneList: d.gauge_update(text=bone,percent=int(boneCount*100/len(boneList)),update_text=True) hashes.write(sha256sum(bone) + " " + bone + "\n") boneCount += 1 d.gauge_update(text="Complete!",percent=100,update_text=True) time.sleep(1) hashes.close() def finish(allSuccessful): bagAndTag() msg = "" if(allSuccessful): msg = "All modules finished execution" else: msg = "Some or all modules failed execution. Please check the logs." code, tag = d.menu(msg, choices=[("Poweroff", "Shutdown the computer"), ("Reboot", "Reboot the computer"), ("Restart", "Run Gravedigger again")]) if (code == d.OK): if tag == "Poweroff": sp.call("poweroff", shell=True) elif tag == "Reboot": sp.call("reboot", shell=True) logger.msgLog("Finished excavation of tomb: " + tomb.getPath(),"digger",logger.TYPE_INFO) exit(0) def showContinueDialog(d,msg): code, tag = d.menu(msg + "\n\nDo you want to retry module configuration or continue execution without the module?", choices=[("retry", "Retry configuration"), ("continue", "Continue without module")]) if code != d.OK: return "abort" return tag if __name__ == "__main__": locale.setlocale(locale.LC_ALL, '') d = Dialog(dialog="dialog",autowidgetsize=True) d.set_background_title("Gravedigger") moduleList = getModules() code,value = d.inputbox("Input computer's name") if code == d.OK: tomb._MACHINE_NAME= value logger.logSystems.append(ConsoleLogger()) logger.logSystems.append(FileLogger(tomb.getPath() + "log.txt")) logger.msgLog("Beginning excavation of tomb: " + tomb.getPath(),"digger",logger.TYPE_INFO) d.set_background_title("Gravedigger - " + value) code, tags = d.checklist("Select modules to execute", choices=moduleList + [("all","Execute all modules",False)], title="Module selection") if code == d.OK: runlist = [] if "all" in tags: logger.msgLog("Selected modules: " + repr(moduleList), "digger", logger.TYPE_INFO) for module in moduleList: while True: try: mod = prepareModule(module[0]) runlist.append(mod) break except Exception as e: msg = "Exception raised while preparing module \"" + module[0] + "\": " + str(e) logger.msgLog(msg, module[0], logger.TYPE_ERROR) #logger.msgLog(traceback.format_exc(), module[0], logger.TYPE_ERROR) ans = showContinueDialog(d,msg) if ans == "abort": finish(False) elif ans == "continue": logger.msgLog("Skipping module " + module[0], "digger", logger.TYPE_INFO) break else: logger.msgLog("Selected modules: " + repr(tags), "digger", logger.TYPE_INFO) for tag in tags: while True: try: mod = prepareModule(tag) runlist.append(mod) break except Exception as e: msg = "Exception raised while preparing module \"" + tag + "\": " + str(e) logger.msgLog(msg,tag,logger.TYPE_ERROR) #logger.msgLog(traceback.format_exc(), tag logger.TYPE_ERROR) ans = showContinueDialog(d, msg) if ans == "abort": finish(False) elif ans == "continue": logger.msgLog("Skipping module " + tag, "digger", logger.TYPE_INFO) break sp.call('clear', shell=True) for module in runlist: logger.msgLog("Running " + module.name + "...", "digger", logger.TYPE_INFO) try: #Battle routine, set! module.execute() except Exception as e: logger.msgLog("Exception raised while running \"" + module.name + "\": " + str(e), module.name, logger.TYPE_ERROR) #logger.msgLog(traceback.format_exc(), module.name, logger.TYPE_ERROR) finish(True)