diff --git a/keygen.py b/keygen.py index c9e58c6..5279ac2 100644 --- a/keygen.py +++ b/keygen.py @@ -1,4 +1,5 @@ from cryptography.fernet import Fernet +from cryptography.fernet import InvalidToken from pathlib import Path import settings import string @@ -6,14 +7,16 @@ import secrets import sys import os -# Check if the script is ran as root + +# Check if the script was run as root if os.geteuid() != 0: exit("Root privileges are necessary to run this script.\nPlease try again as root or using `sudo`.") -# Check if encryption key exists + +# Check if encryption key already exists enckey = Path(settings.ENCKEY_PATH) if enckey.is_file(): - print("Encryption key found.") + print("Encryption key found.") else: print("Encryption key not found.") print("Generating key...") @@ -22,32 +25,18 @@ else: key_file.write(key) print("Encryption key generated and stored in secret.key.") + # Load encryption key def load_key(): - return open(settings.ENCKEY_PATH, "rb").read() + with open(settings.ENCKEY_PATH, "rb") as kf: + kdata = kf.read() + return kdata -# Set size of string -N = 64 - -# Generating of key -token = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(64)) - -# Decrypt the existing keyfile -key = load_key() -f = Fernet(key) -with open("uploadkeys", "rb") as file: - # read the encrypted data - encrypted_data = file.read() -# decrypt data -decrypted_data = f.decrypt(encrypted_data) -# write the original file -with open("uploadkeys", "wb") as file: - file.write(decrypted_data) # Encrypting and storing of key def encrypt_key(message): key = load_key() - f = Fernet(key) + keyf = Fernet(key) with open('uploadkeys', 'a+') as uploadkeys: print(str(token), file=uploadkeys) @@ -55,11 +44,59 @@ def encrypt_key(message): with open("uploadkeys", "rb") as keyfile: keyfile_data = keyfile.read() - encrypted_data = f.encrypt(keyfile_data) + encrypted_data = keyf.encrypt(keyfile_data) with open("uploadkeys", "wb") as keyfile: keyfile.write(encrypted_data) -# Print result on display and call encrypt_key -print("Your new token is: " + str(token)) -encrypt_key(str(token)) \ No newline at end of file + +def ask_yn(msg): + resps = {"y": True, "n": False} + ask = True + while ask: + proceedraw = input(msg) + if proceedraw.lower() in resps.keys(): + proceed = resps[proceedraw] + ask = False + else: + print("Invalid response.") + return proceed + + +N = 64 # Size of token + +# Generate key +token = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(N)) + +# 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() + + 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] ") + + 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