split-secret/scripts/main.py

import argparse
from classes.Encryption import Encryption
from classes.Cleanup import Cleanup
from classes.Decryption import Decryption
from getpass import getpass
from classes.AbstractSplittedSecret import AbstractSplittedSecret
import traceback

cleanup = Cleanup()

def clean_exit():
    print("Cleaning up.")
    cleanup.cleanupFiles(AbstractSplittedSecret.TYPE_DECRYPTED)
    print("Leaving program. Goodby :)")
    exit()
    pass
try:
    if __name__ == '__main__':
        parser = argparse.ArgumentParser()
        parser.add_argument('--mode',type=str, dest='mode',required=True,choices=['cleanup','encrypt','decrypt'])
        parser.add_argument('--amount',type=int, dest='amount_of_secret_holders',required=False,choices=AbstractSplittedSecret.getCoSecretHoldersRange())
        parser.add_argument('--quota', type=int, dest='decryption_quota', choices=range(1,101),required=False)
        parser.add_argument('--master-password',type=str, dest='master_password',required=False)
        parser.add_argument('--user-password',type=str, dest='user_password',required=False)
        parser.add_argument('--user',type=int, dest='user',choices=AbstractSplittedSecret.getSecretHoldersRange(),required=False)
        parser.add_argument('--add-user-information',type=bool, dest='add_user_information', default=False, required=False, action=argparse.BooleanOptionalAction)
        args = parser.parse_args()

        print("Splitted Secret Interface started.")
        print("Selected Mode: " + args.mode)
        
        if args.mode == 'cleanup':   
            if args.user is None: 
                print("Delete all files.")
                cleanup.deleteAll()
                clean_exit()
            print("Delete files for user <<" + str(args.user) + ">>");
            cleanup.cleanupForUser(args.user)
            clean_exit()
            
        if args.mode == 'decrypt':
            decrypt = Decryption()
            if args.master_password is None:
                if args.user is None: 
                    print("Type in the user id:")
                    decrypt.initializeUser(input())
                else:
                    decrypt.initializeUser(args.user)
                if args.user_password is None:
                    while True:
                        print("Enter the user password:")
                        decrypt.setUserPassword(getpass())
                        print("Decrypting User File...")
                        try:
                            decrypt.initializeUserDataDecryption();
                            break;
                        except Exception as error:
                            print("An error occured. Propably you typed in a wrong password :( The error is: " + str(error))
                else:
                    decrypt.setUserPassword(args.user_password)
                    print("Decrypting User File...")
                    try:
                        decrypt.initializeUserDataDecryption();
                    except Exception as error:
                        print("An error occured. Propably you passed a wrong password :( The error is: " + str(error))
                        clean_exit()
                print("Contact the following persons and tell them that you need help to encrypt the data: \n")
                for contact_id in decrypt.user_data['contacts']:
                    print("user_id: " + contact_id)
                    for label in decrypt.user_data['contacts'][contact_id]:
                        print(label + ": " + decrypt.user_data['contacts'][contact_id][label])
                    print()
                while True:
                    decrypt.resetDecrypterIds()
                    try:
                        person_counter = 1
                        while person_counter <= decrypt.getNeededCoDecryptersAmount():
                            print("The following user id's are in the decryption list: " + str(decrypt.getDecrypterIds()))
                            print("You need at least <<" + str(decrypt.getNeededCoDecryptersAmount()) +">> other person to decrypt the secret.")
                            print("Type in the user id of another encrypter:")
                            decrypt.addDecrypterId(int(input()))
                            person_counter += 1
                        break
                    except Exception as error:
                        print("The following error occured <<" + str(error) + ">> :( \n Try again :)")
                print("\nYour data is:\n")
                print("FOR PASSWORD GROUP: " + decrypt.getDecryptersGroupName()) 
                print("FOR USER ID: "  + decrypt.getUserId())
                print("PASSWORD SHARE IS: " + decrypt.getPasswordShare() + "\n")
                while True:
                    try:
                        decrypt.resetPasswordShare()
                        co_decrypter_ids = decrypt.getCoDecrypterIds()
                        for co_decrypter_id in decrypt.getCoDecrypterIds():
                            print("Execute this script for user: " + str(co_decrypter_id) + ".")
                            print("Type in the password share.\n")
                            print("\nFOR PASSWORD GROUP: " + decrypt.getDecryptersGroupName()) 
                            print("FOR USER: " + str(co_decrypter_id)) 
                            print("PASSWORD SHARE IS: ")
                            decrypt.addPasswordShare(co_decrypter_id, input())
                        print("\nTHE GROUP PASSWORD IS: " + decrypt.getGroupPassword())
                        print("\nDecrypting group password file.\n")
                        decrypt.initializeGroupDataEncryption()
                        print("THE MASTER PASSWORD IS: " + decrypt.getMasterPassword())
                        break;
                    except:
                        print("An unexpected error occured: \n" + traceback.format_exc())
                clean_exit()
            print("Decrypting accumulated file...")
            decrypt.setUserPassword(args.master_password)
            decrypt.decryptAccumulatedFile()
            clean_exit()
        
        if args.mode == 'encrypt':
            if args.master_password is None:
                print("Enter the master password:")
                master_password = getpass()
            else:
                master_password = args.master_password
            encrypt = Encryption(args.amount_of_secret_holders, args.decryption_quota, master_password)
            if args.add_user_information is not None:
                for user_id in encrypt.user_mapped_data:
                    for label in ['name','phone','email','address']:
                        print("Enter attribut <<" + label + ">> for user <<" + user_id+ ">>:" )
                        encrypt.addInformationToUser(user_id, label, str(input()))
            encrypt.compileData()
            encrypt.encrypt()
            clean_exit()
except KeyboardInterrupt:
    print("Program interrupted by user.")
except:
    print("An unexpected error occured: \n" + traceback.format_exc())
clean_exit()