Imanol-Mikel Barba Sabariego / gravedigger · Files

#! /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)