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!

if_verbose = False

def verbose(string):
    if if_verbose: print(string)

# 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}")

# Encrypting data
def encrypt_data_neo(unencrypted_name, password, encrypted_name=None, use_key=None, keyfile_name=None, delete_original=None):

    verbose(f"original file: {unencrypted_name}")

    if not encrypted_name:
        encrypted_name="encrypted"
    
    file_chosen_name=encrypted_name
    verbose(f"new file: {encrypted_name}")

    if use_key:
        verbose("Using key")
    else:
        verbose("not using key")
    
    verbose(f"password: {password}")
    password = password.encode()
    verbose(f"password encoded: {password}")

    salt = os.urandom(16)

    verbose(f"salt: {salt}")

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=480000,
    )

    key = base64.urlsafe_b64encode(kdf.derive(password))

    verbose(f"key: {key}")

    f = Fernet(key)

    verbose(f"f: {f}")

    try:
        with open(unencrypted_name, 'rb') as file:
            original_data = file.read()
        
        verbose(f"{unencrypted_name} has been read")

        encrypted_data = f.encrypt(original_data)

        verbose(f"data has been encrypted.")

        if not use_key:
            with open(f"{encrypted_name}", 'wb') as file:
                # Write the 16-byte salt first, then the encrypted data
                file.write(salt)
                file.write(encrypted_data)
            verbose(f"data was written to {encrypted_name} without key.")

        elif use_key:
            if not keyfile_name:
                keyfile_name = "key.p7c_key"
            
            verbose(f"keyfile_name: {keyfile_name}")

            with open(f"{keyfile_name}", 'wb') as key:
                key.write(salt)

            verbose("salt has been written.")

            with open(f"{encrypted_name}", 'wb') as file:
                file.write(encrypted_data)

            verbose("encrypted data has been written.")
        else:
            verbose("🔥🔥🔥 THIS IS AN ERROR 🔥🔥🔥")
            print(f"This is P7MJ. I am speechless. use_key: {use_key}")

        if delete_original:
            os.remove(f"{unencrypted_name}")
            verbose("Original file has been removed.")

    except Exception as e:
        verbose("🔥🔥🔥 THIS IS AN ERROR 🔥🔥🔥")
        print(f"\nAn error has occured.\n {e}\n")

# 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}.")

# Decrypting data
def decrypt_data_neo(encrypted_name, password, unencrypted_name=None, use_key=None, keyfile_name=None, delete_original=None):

    choose_file = encrypted_name

    if not unencrypted_name:
        unencrypted_name = "unencrypted"

    password = password.encode()

    if use_key and not keyfile_name:
        keyfile_name = "key.p7c_key"

    # 2. Open the encrypted file and extract the salt + data
    if use_key:
        with open(f"{keyfile_name}", 'rb') as file:
            file_salt = file.read(16)
        with open(choose_file, 'rb') as file:
            encrypted_data = file.read()
    
    elif not use_key:
        with open(encrypted_name, '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:
        verbose("🔥🔥🔥 THIS IS AN ERROR 🔥🔥🔥 flames are burning and so is my brain")
        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))

    verbose(f"key: {key}")

    f = Fernet(key)

    verbose(f"f: {f}")

    try:
        decrypted_data = f.decrypt(encrypted_data)
        verbose("data unencrypted.")
        
        with open(f'{unencrypted_name}', 'wb') as file:
            file.write(decrypted_data)
        verbose("unencrypted data written to disk.")
        if delete_original:
            os.remove(encrypted_name)
            verbose("Encrypted file removed.")
            if use_key:
                os.remove(keyfile_name)
            verbose("Key removed.")
        
    except Exception as e:
        print(f"\nError. Could not decrypt. It is likely that your password is wrong.\nError: {e}.\n")

# Main function
def old_main():
    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!!!")

if __name__ == "__main__":
    try:
        while True:
            inputs = input("[PCS] ")
            list_inputs = inputs.split()

            if not list_inputs:
                continue

            for i, param in enumerate(list_inputs):
                if param == "None" or param == "0" or param == "No" or param == "no" or param == "False":
                    list_inputs[i] = None
                if param == "True" or param == "1" or param == "Yes" or param == "yes" or param == "Exists" or param == "exists":
                    list_inputs[i] = True
            
            if list_inputs[0] == "help":
                print("encrypt unencrypted_name encrypted_name=None password use_key=None keyfile_name=None delete_original=None")
                #              1               2                    3        4            5                 6
                print("decrypt encrypted_name unencrypted_name=None password use_key=None keyfile_name=None delete_original=None")
                #              1              2                     3        4            5                 6
                print("\nDue to this program being able to encrypt absolutely ANYTHING, It is recommended to name your encrypted archives after your original file.")
                print("E.g. archive.zip should be named archive.zip.p7c_enc.")
                print("This ensures that you remember the extension.")
                print("\nFor compatibility with PCS A-2-i Dividend, the keyfile must be named key.p7c_key. Additionally, append .p7c_enc to all encrypted archive names.")
                print("For compatability with PCS A-1-i Proto, follow A-2-i standards, but do not use any form of keyfile.")
            
            elif list_inputs[0] == "credits":
                print("PCS A-3-i Terminal by P7MJ.")
            
            elif list_inputs[0] == "exit":
                sys.exit(0)

            elif list_inputs[0] == "encrypt":
                if len(list_inputs) == 7:
                    # unencrypted_name, password, encrypted_name=None, use_key=None, keyfile_name=None, delete_original=None
                    encrypt_data_neo(list_inputs[1], list_inputs[3], encrypted_name=list_inputs[2], use_key=list_inputs[4], keyfile_name=list_inputs[5], delete_original=list_inputs[6])
                else:
                    print("Invalid use of encrypt: length of parameters incorrect.")
                    print('If you wish to leave a parameter empty, type "None", "0", "No", "no", or "False".')
                    print('Type "help" for more information.')
            
            elif list_inputs[0] == "decrypt":
                if len(list_inputs) == 7:
                    # encrypted_name, password, unencrypted_name=None, use_key=None, keyfile_name=None, delete_original=None
                    decrypt_data_neo(list_inputs[1], list_inputs[3], unencrypted_name=list_inputs[2], use_key=list_inputs[4], keyfile_name=list_inputs[5], delete_original=list_inputs[6])
                else:
                    print("Invalid use of decrypt: length of parameters incorrect.")
                    print('If you wish to leave a parameter empty, type "None", "0", "No", "no", or "False".')
                    print('Type "help" for more information.')
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"{e}")