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

# EXPERIMENTAL

def main(args):
    # 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":
            break
        else:
            print("Not an option!!!")