digger.py 7.4 KB
#! /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
import bonetypes
import xml.etree.ElementTree
import binascii

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:
            logger.msgLog("Exception raised while importing " + modname + ": " + str(e),"digger",logger.TYPE_ERROR)
    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()
    if os.path.isfile(tombPath + "MANIFEST.XML"):
        etree = xml.etree.ElementTree.parse(tombPath + "MANIFEST.XML")
        xmlroot = etree.getroot()
    else:
        xmlroot = xml.etree.ElementTree.Element("Manifest")
        newElem = xml.etree.ElementTree.Element("Case")
        newElem.text = tomb.__getTomb__()
        xmlroot.insert(0,newElem)
        xmlroot.insert(1,xml.etree.ElementTree.Element("EvidenceList"))
        etree = xml.etree.ElementTree.ElementTree(xmlroot)

    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({"path" : subroot + "/" + filename, "type": bonetypes.definitions[dir]})

    d.gauge_start(title="Hashing all collected artifacts...",width=60,height=10)
    boneCount = 0
    evidenceList = xmlroot.find("EvidenceList")
    for bone in boneList:
        relpath = bone["path"][len(tombPath):]
        if evidenceList.find("./Evidence[@path='" + relpath + "']") != None:
            continue
        hash = sha256sum(bone["path"])
        date = str(int(os.path.getmtime(bone["path"])*1000))
        evidenceList.insert(len(evidenceList),xml.etree.ElementTree.Element("Evidence", hash=hash,path=relpath,type=bone["type"],date=date))
        d.gauge_update(text=bone["path"],percent=int(boneCount*100/len(boneList)),update_text=True)
        boneCount += 1
    d.gauge_update(text="Complete!",percent=100,update_text=True)
    time.sleep(1)
    etree.write(tombPath + "MANIFEST.XML")

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")
    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)
        moduleList = getModules()
        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)