From 9f2c7c2b88094527dc8c8cc5c6a96fba7518b2a0 Mon Sep 17 00:00:00 2001 From: BBaoVanC Date: Thu, 3 Sep 2020 20:20:22 -0500 Subject: [PATCH] Clean up and reorganize keygen.py This makes it easier to add new features - Clearer variable names - Split up large functions into smaller ones for each action - Help make it easier to add new features --- keygen.py | 146 +++++++++++++++++++++++++++++++----------------------- 1 file changed, 83 insertions(+), 63 deletions(-) diff --git a/keygen.py b/keygen.py index 1ec3643..bb3874e 100644 --- a/keygen.py +++ b/keygen.py @@ -8,41 +8,63 @@ import sys import os -# Check if encryption key already exists -enckey = Path(settings.ENCKEY_PATH) -if enckey.is_file(): - print("Encryption key found.") -else: - print("Encryption key not found.") - print("Generating key...") - key = Fernet.generate_key() - with open(settings.ENCKEY_PATH, "wb") as key_file: - key_file.write(key) - print("Encryption key generated and stored in secret.key.") - - -# Load encryption key -def load_key(): - with open(settings.ENCKEY_PATH, "rb") as kf: - kdata = kf.read() - return kdata +# Load secret +def load_secret(): + with open(settings.ENCKEY_PATH, "rb") as sf: + secret = sf.read() + return secret # Encrypting and storing of key -def encrypt_key(message): - key = load_key() - keyf = Fernet(key) +def append_uploadkey(akey): + with open('uploadkeys', 'a+') as uploadkeysf: + print(str(akey), file=uploadkeysf) - with open('uploadkeys', 'a+') as uploadkeys: - print(str(token), file=uploadkeys) - with open("uploadkeys", "rb") as keyfile: - keyfile_data = keyfile.read() +def decrypt_uploadkeys(): + with open("uploadkeys", "rb") as uploadkeysf: + uploadkeys_data = uploadkeysf.read() - encrypted_data = keyf.encrypt(keyfile_data) + try: + secret = load_secret() + secretf = Fernet(secret) + decrypted_data = secretf.decrypt(uploadkeys_data) # decrypt data + with open("uploadkeys", "wb") as ukf: + ukf.write(decrypted_data) # write the original file + print("Done decrypting") # debug + return True + except InvalidToken: + print("InvalidToken") # debug + print("The encrypted key data is invalid and cannot be read.") + print("It may be necessary to clear the file entirely, which will invalidate all tokens.") + proceed = ask_yn("Do you wish to proceed to clearing the uploadkeys file? [y/n] ") - with open("uploadkeys", "wb") as keyfile: - keyfile.write(encrypted_data) + if proceed: + print("Proceed1") + os.remove("uploadkeys") + print("Removed uploadkeys file.") + proceed2 = ask_yn("Would you like to continue and generate a new key? [y/n] ") + if not proceed2: + print("not proceed2") + return False + else: + print("proceed2") + return True + else: + print("not Proceed1") + return False + + +def encrypt_uploadkeys(): + with open("uploadkeys", "rb") as uploadkeysf: + uploadkeys_data = uploadkeysf.read() + + secret = load_secret() + secretf = Fernet(secret) + encrypted_data = secretf.encrypt(uploadkeys_data) + + with open("uploadkeys", "wb") as uploadkeysf: + uploadkeysf.write(encrypted_data) def ask_yn(msg): @@ -58,46 +80,44 @@ def ask_yn(msg): return proceed -start = ask_yn("Have you run this program as the correct user (for example, nginx uses www-data)? [y/n] ") -if not start: - print("Please run this as the correct user with: sudo su [user] -s /bin/sh -c 'python3 keygen/py'") - +# Check if encryption secret already exists +if Path(settings.ENCKEY_PATH).is_file(): + print("Encryption secret found.") else: + print("Encryption secret not found.") + print("Generating secret...") + newsecret = Fernet.generate_key() + with open(settings.ENCKEY_PATH, "wb") as secret_file: + secret_file.write(newsecret) + print("Encryption secret generated and stored in {0}".format(settings.ENCKEY_PATH)) - N = 64 # Size of token - # Generate key - token = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(N)) +if __name__ == "__main__": + start = ask_yn("Have you run this program as the correct user (for example, nginx uses www-data)? [y/n] ") + if not start: + print("Please run this as the correct user with: sudo su [user] -s /bin/sh -c 'python3 keygen.py'") - # Decrypt the existing keyfile - key = load_key() - keyf = Fernet(key) - - genkey = True - uploadkeysp = Path("uploadkeys") - if not uploadkeysp.is_file(): - uploadkeysp.touch() else: - with open("uploadkeys", "rb") as ukf: - # read the encrypted data - encrypted_data = ukf.read() + # Decrypt the existing keyfile + secret = load_secret() + keyf = Fernet(secret) - try: - decrypted_data = keyf.decrypt(encrypted_data) # decrypt data - with open("uploadkeys", "wb") as ukf: - ukf.write(decrypted_data) # write the original file - except InvalidToken: - print("The encrypted key data is invalid and cannot be read.") - print("It may be necessary to clear the file entirely, which will invalidate all tokens.") - proceed = ask_yn("Do you wish to proceed to clearing the uploadkeys file? [y/n] ") + genkey = True + uploadkeysp = Path("uploadkeys") + if not uploadkeysp.is_file(): + uploadkeysp.touch() + else: + with open("uploadkeys", "rb") as ukf: + # read the encrypted data + encrypted_data = ukf.read() - if proceed: - os.remove("uploadkeys") - print("Removed uploadkeys file.") - proceed2 = ask_yn("Would you like to continue and generate a new token? [y/n] ") - if not proceed2: - genkey = False - if genkey: - print("Your new token is: " + str(token)) # Print token - encrypt_key(str(token)) # Encrypt the key and save + if genkey: + if decrypt_uploadkeys(): # Decrypt the file + N = 64 # Size of key + key = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(N)) + print("Your new key is: " + str(key)) # Print key + append_uploadkey(key) # Save the new key to file unencrypted + encrypt_uploadkeys() # Encrypt the uploadkeys file + else: + print("Exiting.")