import os
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

# Function to encrypt a file using a salt
def encrypt_file(input_file, output_file=None):
    key = os.urandom(16)
    iv = os.urandom(16)

    cipher = Cipher(algorithms.AES(key), modes.CFB(iv))
    encryptor = cipher.encryptor()

    encrypted_content = b""

    if output_file is not None:
        with open(input_file, 'rb') as infile, open(output_file, 'wb') as outfile:
            # Write the IV to the output file first
            outfile.write(iv)
            encrypted_content += iv

            while chunk := infile.read(2048):
                ciphertext = encryptor.update(chunk)
                outfile.write(ciphertext)
                encrypted_content += ciphertext

            # Finalize encryption
            final_chunk = encryptor.finalize()
            outfile.write(final_chunk)
            encrypted_content += final_chunk

    return key, encrypted_content


# Function to decrypt a file
def decrypt_file(key, input_file, output_file=None):
    plaintext_content = b""

    with open(input_file, 'rb') as infile:
        # Read the IV from the input file
        iv = infile.read(16)
        cipher = Cipher(algorithms.AES(key), modes.CFB(iv))
        decryptor = cipher.decryptor()

        if output_file is not None:
            with open(output_file, 'wb') as outfile:
                while chunk := infile.read(2048):
                    plaintext = decryptor.update(chunk)
                    outfile.write(plaintext)
                    plaintext_content += plaintext

                # Finalize decryption
                final_chunk = decryptor.finalize()
                outfile.write(final_chunk)
                plaintext_content += final_chunk
        else:
            while chunk := infile.read(2048):
                plaintext = decryptor.update(chunk)
                plaintext_content += plaintext

            # Finalize decryption
            plaintext_content += decryptor.finalize()

    return plaintext_content