Main.py

from classes.MalwareWatcher import MalwareWatcher
from classes.MalwareReverter import MalwareReverter
from classes.utils import isAdmin, printHeader
import argparse

def initArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--changeFile",
                        help="The file that has been created by analyzing. This is used to revert "
                             "some changes")
    args = parser.parse_args()

    return args

def main(changeFile):
    if not isAdmin():
        input("[!] Run this as admin.")
        exit()

    printHeader()

    if changeFile:
        Mr = MalwareReverter()
        Mr.revertChanges(changeFile)

    else:
        Mw = MalwareWatcher()
        Mw.run()

    input("\nPress ENTER to close...")


if __name__ == '__main__':
    args = initArgs()
    main(args.changeFile)