digger.py 6.05 KB
#! /usr/bin/env python3

import sys

sys.path.append('modules')
sys.path.append('vars')
sys.path.append('logger')

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+')

    for root, dirnames, filenames in os.walk(tombPath):
        for dir in dirnames:
            print(dir)
            for subroot, subdirnames, subfilenames in os.walk(tombPath + dir):
                print(subfilenames)
                for filename in subfilenames:
                    hashes.write(sha256sum(subroot + "/" + filename) + " " + subroot + "/" + filename + "\n")

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:
                    module.run()
                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)