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!!!")