diff --git a/PCS-A-2-i Dividend.py b/PCS-A-2-i Dividend.py new file mode 100644 index 0000000..5a0d9a0 --- /dev/null +++ b/PCS-A-2-i Dividend.py @@ -0,0 +1,169 @@ +import os +import getpass +import sys +from cryptography.fernet import Fernet +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +from pathlib import Path +import base64 + +# TODO Add option to write and read keyfile, where there is an option to store the salt! + +# Encrypting data +def encrypt_data(): + + # Pesterish file asker + while True: + file_name = input("\nEnter the name of a zip file, without the extension that is in the same directory as this script. \nE.g. for encrypt.zip you would type \"encrypt\". Name > ") + ".zip" + if Path(file_name).is_file(): + break + else: + print("This zip file does not exist!") + + file_chosen_name = input("\nMake a name for your encrypted file, e.g. encrypted (press enter for default)\nName > ") + if file_chosen_name == "": + file_chosen_name = "encrypted" + + keyfile = input("\nDo you want enable keyfile encryption?\nIf enabled, decryption is only possible with the file.[y]/[N] ") + if keyfile == "y" or keyfile == "Y": + print("Keyfile enabled.") + enable_keyfile = True + else: + print("Keyfile disabled.") + enable_keyfile = False + + # Pestering password asker until you confirm it + while True: + password = getpass.getpass("\nCreating password > ") + confirm = input(f"Did you type the correct password, and remember it? [Y]/[n] ") + if confirm == "Y" or confirm == "y" or confirm == "": + password = password.encode() + break + + # 1. Setup Salt + salt = os.urandom(16) + + # 2. Derive Key + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=32, + salt=salt, + iterations=480000, + ) + key = base64.urlsafe_b64encode(kdf.derive(password)) + f = Fernet(key) + + print("Encrypting...", end = " ") + try: + # 3. Encrypt and save SALT + DATA together + with open(file_name, 'rb') as file: + original_data = file.read() + + encrypted_data = f.encrypt(original_data) + + if not enable_keyfile: + with open(f"{file_chosen_name}.p7c_enc", 'wb') as file: + # Write the 16-byte salt first, then the encrypted data + file.write(salt) + file.write(encrypted_data) + + elif enable_keyfile: + with open(f"key.p7c_key", 'wb') as key: + key.write(salt) + + with open(f"{file_chosen_name}.p7c_enc", 'wb') as file: + file.write(encrypted_data) + else: + print("What the f** did you do something to enable_keyfile?") + + print("Encrypted! ") + delete_original = input("Delete the original zip file for security? [Y]/[n] ") + if delete_original == "Y" or delete_original == "y" or delete_original == "": + os.remove(f"{file_name}") + print("Original file removed.") + except Exception as e: + print(f"Error! {e}") + +# Decrypting data +def decrypt_data(): + # 1. Get the password from the user + while True: + choose_file = input("\nName of the p7c_enc file without extension, e.g. \"encrypted\"\nName > ") + ".p7c_enc" + if Path(choose_file).is_file(): + break + else: + print("This p7c_enc file does not exist!") + + key_exists = input("Do you have a p7c_key keyfile? [y]/[N] ") + if key_exists == "y" or key_exists == "Y": + while True: + input("Double check if your key is in the same direcory as this code and your encrypted file.\nIt must be named key.p7c_key. Press ENTER to check.") + if Path("key.p7c_key").is_file(): + break + else: + print("Does not exist!") + there_is_a_key = True + else: + there_is_a_key = False + + password = getpass.getpass("Password > ").encode() + + # 2. Open the encrypted file and extract the salt + data + if there_is_a_key: + with open("key.p7c_key", 'rb') as file: + file_salt = file.read(16) + with open(choose_file, 'rb') as file: + encrypted_data = file.read() + + elif not there_is_a_key: + with open(choose_file, 'rb') as file: + # Read exactly 16 bytes for the salt + file_salt = file.read(16) + # Read everything else as the encrypted data + encrypted_data = file.read() + else: + print("WHAT DID YOU DO TO THERE_IS_A_KEY?!") + + # 3. Re-derive the EXACT same key using that salt + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=32, + salt=file_salt, + iterations=480000, + ) + key = base64.urlsafe_b64encode(kdf.derive(password)) + f = Fernet(key) + + # 4. Decrypt and save the original file + print("Decrypting...", end = " ") + try: + decrypted_data = f.decrypt(encrypted_data) + + with open('unencrypted.zip', 'wb') as file: + file.write(decrypted_data) + print("Success.") + delete_encrypted = input("\nDelete the encrypted file (and key, if exists)? [Y]/[n] ") + if delete_encrypted == "Y" or delete_encrypted == "y" or delete_encrypted == "": + os.remove(choose_file) + if there_is_a_key: + os.remove("key.p7c_key") + print("Key removed.") + print("Encrypted file removed.") + + + except Exception as e: + print(f"Could not decrypt!.\nThis might be due to a wrong file or corrupted data?\nDetailed info:\n{e}.") + +# Main function + +while True: + print("\nPCS - P7MJ's enCryption System | V A-2-i Dividend | P7MJ") + choice = input("[1] Encrypt [2] Decrypt [x] Exit > ") + if choice == "1": + encrypt_data() + elif choice == "2": + decrypt_data() + elif choice == "x": + sys.exit(0) + else: + print("Not an option!!!")