From 2e5b085cd9e00fe8a4bb58d751a2576b8f8b35cf Mon Sep 17 00:00:00 2001 From: Kai-Denzel_jane Date: Wed, 17 Jan 2024 23:28:33 +0800 Subject: [PATCH] Fixed some bugs and preparing for file encryption and decryption --- .env | 4 +- Encryption/Fernet (Local-Files)/Fernet.py | 366 ++++++++++++++++++---- 2 files changed, 299 insertions(+), 71 deletions(-) diff --git a/.env b/.env index a0c0008..37a8039 100644 --- a/.env +++ b/.env @@ -1,3 +1,3 @@ KEY='' -PASSWORD='Alpine' -KEY_BACKUP='' \ No newline at end of file +PASSWORD='YWxwaW5l' +KEY_BACKUP='' diff --git a/Encryption/Fernet (Local-Files)/Fernet.py b/Encryption/Fernet (Local-Files)/Fernet.py index d80c754..a5918e4 100644 --- a/Encryption/Fernet (Local-Files)/Fernet.py +++ b/Encryption/Fernet (Local-Files)/Fernet.py @@ -3,19 +3,44 @@ from dotenv import load_dotenv import dotenv from cryptography.fernet import Fernet +import base64 -load_dotenv() +def load_env_variables(): + load_dotenv() + global KEY, PASSWORD_E, PASSWORD_D, KEY_BACKUP + KEY = os.getenv("KEY") + PASSWORD_E = os.getenv("PASSWORD") + PASSWORD_D = base64.b64decode(PASSWORD_E).decode("utf-8", "strict") + KEY_BACKUP = os.getenv("KEY_BACKUP") -KEY = os.getenv("KEY") -PASSWORD = os.getenv("PASSWORD") -KEY_BACKUP = os.getenv("KEY_BACKUP") +def validate_input(prompt, expected_type, error_message): + while True: + user_input = input(prompt) + try: + user_input = expected_type(user_input) + return user_input + except ValueError: + print(error_message) +load_env_variables() + # Asks users for their password def pass_check(): + """ + Asks the user for password and checks if it is correct + + Args: + None + + Returns: + bool: validity of password + """ + + user_input = input("Enter password: ") - if user_input == PASSWORD: + if user_input == PASSWORD_D: validity = True return(validity) @@ -24,41 +49,104 @@ def pass_check(): print("Try again.") pass_check() -# Welcome function + def welcome(): - print("1. Encrypt") - print("2. Decrypt") - print("3. Generate new key") - print("4. Input key and use to encrypt / decrypt") - print("5. Print out your current key") - print("6. Set password (Advanced Users only!!!)") - print("7. Reset Password and Key") - print("8. Manage Keys") + """ + Prints out the welcome screen and asks the user for their choice + + Args: + None - choice = int(input("Input number: ")) - return choice + Returns: + int: users menu choice + + """ + while True: + print("1. Encrypt") + print("2. Decrypt") + print("3. Generate new key") + print("4. Input key and use to encrypt / decrypt") + print("5. Print out your current key") + print("6. Set password (Advanced Users only!!!)") + print("7. Reset Password and Key") + print("8. Manage Keys") + + choice = validate_input("Input a number: ", int, "Please input a number") + + try : + choice = int(choice) + return(choice) + + except ValueError: + print("Please input a number") + + # Encrypt function def encrypt_func(): - text = str(input("Input text: ")) + """ + Encrypts the text the user inputs + + Args: + None + + Returns: + str: encrypted text + + """ + + text = input("Input text: ") + + try: + text = str(text) + + except ValueError: + print("Please input text") + encrypt_func() + data = bytes(text, encoding="utf-8") - cipher = Fernet(KEY) - encrypted = cipher.encrypt(data) + encryption_tool = Fernet(KEY) + encrypted = encryption_tool.encrypt(data) return(str(encrypted)) # Decrypt function def decrypt_func(): + """ + Decrypts the text the user inputs - if pass_check() == True: + Args: + None - data = str(input("Input encrypted text: ")) - - cipher = Fernet(KEY) - decrypted = cipher.decrypt(data) - return(str(decrypted)) + Returns: + str: decrypted text + """ + + if pass_check() == True: + data = input("Input encrypted text: ") + + # Remove the "b'" and "'" from the input if they exist + if data.startswith("b'") and data.endswith("'"): + data = data[2:-1] + + # Try to decrypt the data + try: + encryption_tool = Fernet(KEY) + decrypted = encryption_tool.decrypt(data.encode()) + except: + print("Decryption failed.") + return + + # Try to decode the decrypted data + try: + decrypted = decrypted.decode("utf-8", "strict") + except: + print("Decoding failed.") + return + + return decrypted else: pass_check() @@ -66,18 +154,43 @@ def decrypt_func(): # Key generation def key_gen(): + """ + Generates a new key and saves it to .env file + + Args: + None + + Returns: + None + + """ + if pass_check() == True: key = Fernet.generate_key() key_str = key.decode("utf-8", "strict") dotenv.set_key(".env", "KEY", key_str) + + print("Key generated") + else: pass_check() # Import custom keys so you can share encrypted messages with others def custom_key(): + """ + Imports a custom key and saves it to .env file + + Args: + None + + Returns: + None + + """ + if pass_check() == True: key = input("Input key: ") @@ -87,11 +200,18 @@ def custom_key(): else: pass_check() +def key_print(): - + """ + Prints out the current key -# Prints out key, you can also read it directly from "thekey.key" file -def key_print(): + Args: + None + + Returns: + None + + """ if pass_check() == True: print(KEY) @@ -100,14 +220,32 @@ def key_print(): # Changes password def set_pass(): + + """ + Changes the password + + Args: + None + + Returns: + None + + """ print("Warning for security reasons this will reset your key as well so you cant access someone elses encryptions by resetting password.") key_gen() - password = (input("Input new Password: ")) + password_d = (input("Input new Password: ")) + double_check = (input("Input new Password again: ")) + + if password_d != double_check: + print("Passwords do not match") + set_pass() + + password_e = base64.b64encode(bytes(password_d, encoding="utf-8")) print("The script will now exit and you password will be affective on next launch") - dotenv.set_key(".env", "PASSWORD", password) + dotenv.set_key(".env", "PASSWORD", password_e) exit() @@ -115,71 +253,157 @@ def set_pass(): # Resets the password and generates a new key def reset(): + """ + Resets the password and generates a new key + + Args: + None + + Returns: + None + + """ + key_gen() default = "Alpine" dotenv.set_key(".env", "PASSWORD", default) def managekeys(): + """ + Manages keys - print("1. Backup a key") - print("2. Delete backed up key") - print("3. Restore key from backed up key") - print("0. Go back to main menu") - option_key = int(input("Input a number: ")) - - match option_key: - - case 1: - if KEY_BACKUP != '': - + Args: + None + + Returns: + None + """ + while True: + print("1. Backup a key") + print("2. Delete backed up key") + print("3. Restore backed up key as current key") + print("0. Go back to main menu") + option_key = validate_input("Input a number: ", int, "Please input a number") + + match option_key: + case 1: + if KEY_BACKUP != '': + print("Delete this backed up key first") backedupkey = input("Input the key: ") dotenv.set_key(".env", "KEY_BACKUP", backedupkey) - else: - print("Delete this backed up key first") - managekeys() - case 2: - ask = input("Are you sure [y/n]]") + break + case 2: + ask = input("Are you sure [y/n]]") + if ask == "y": + dotenv.set_key(".env", "KEY_BACKUP", '') + break + case 3: + ask = input("Are you sure [y/n]") + if ask == "y": + dotenv.set_key(".env", "KEY", KEY_BACKUP) + break + case 0: + break + +def encrypt_file(): + + """ + Encrypts a file + + Args: + None + + Returns: + None + + """ - if ask == "y": - dotenv.set_key(".env", "KEY_BACKUP", '') - else: - welcome() - case 3: - ask = input("Are you sure [y/n]") + if pass_check() == True: + print(os.curdir) + file = input("Input file name: ") + file = open(file, "rb") + data = file.read() + file.close() + encryption_tool = Fernet(KEY) + encrypted = encryption_tool.encrypt(data) + file = open(file, "wb") + file.write(encrypted) + file.close() + else: + pass_check() - if ask == "y": - dotenv.set_key(".env", "KEY", KEY_BACKUP) - else: - welcome() - case 0: - welcome() +def decrypt_file(): + + """ + Decrypts a file + + Args: + None + + Returns: + None + + """ + + if pass_check() == True: + print(os.curdir) + file = input("Input file name: ") + file = open(file, "rb") + data = file.read() + file.close() + encryption_tool = Fernet(KEY) + decrypted = encryption_tool.decrypt(data) + file = open(file, "wb") + file.write(decrypted) + file.close() + else: + pass_check() # Prompts user if they would like to end the script def end(): - end = input("End [y/n]") - if end == "yes" or end == "y": - exit() + """ + Prompts user if they would like to end the script + + Args: + None + + Returns: + None + + """ + + while True: + end = input("End [y/n]") + + try: + end = str(end) + except ValueError: + print("Please input y or n") + continue + + if end.lower() in ["yes", "y", "no", "n"]: + if end.lower() in ["yes", "y"]: + exit() + elif end.lower() in ["no", "n"]: + return + else: + print("Please input y or n") if KEY == '': print("Either this is your first time running the script or YOU changed you key to '',no worries we are generating a new key for you.") - print("Default password is Alpine you should change it after") + print("Default password is 'alpine' you should change it after") key_gen() # Controls the users choice throughout the script while True: - #.env stuff - load_dotenv() - - KEY = os.getenv("KEY") - PASSWORD = os.getenv("PASSWORD") - KEY_BACKUP = os.getenv("KEY_BACKUP") + load_env_variables() choice = welcome() - match choice: + match choice: # I use match statements because they are easier to read and more efficient than if statements case 1: encrypt = encrypt_func() @@ -199,5 +423,9 @@ def end(): reset() case 8: managekeys() + case 9: + encrypt_file() + case 10: + decrypt_file() end() \ No newline at end of file