diff --git a/PCS-A-1-i-Proto.py b/PCS-A-1-i-Proto.py new file mode 100644 index 0000000..ff1f81a --- /dev/null +++ b/PCS-A-1-i-Proto.py @@ -0,0 +1,130 @@ +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 finish functions +# TODO test this on some unimportant data +# TODO Integrate file choosing in both encrypt and decrypt! + +# 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. E.g. for encrypt.zip you would type\"encrypt\".\nName > ") + ".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" + + # Pestering password asker until you confirm it + while True: + password = getpass.getpass("Creating new password > ") + confirm = input(f"Did you type the correct password, and remember it? [Y]/[n]") + if confirm == "Y" or confirm == "y": + # Confirms with user one last time and reminders + 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) + + 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) + + print("Encrypted! ") + delete_original = input("Delete the original zip file for security? [Y]/[n]") + if delete_original == "Y" or delete_original == "y": + 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!") + + password = getpass.getpass("Password > ").encode() + + # 2. Open the encrypted file and extract the salt + data + 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() + + # 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("Delete the encrypted file? [Y]/[n]") + if delete_encrypted == "Y" or delete_encrypted == "y": + os.remove(choose_file) + 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-1-i Proto | 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!!!")