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 derive a 256-bit key from a password and salt
def derive_key(salt):
	kdf = PBKDF2HMAC(
		algorithm=hashes.SHA256(),
		length=32,
		salt=salt,
		iterations=10000,
		backend=default_backend()
	)
	return kdf.derive(b'')


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

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

	with open(input_file, 'rb') as f:
		plaintext = f.read()

	ciphertext = encryptor.update(plaintext) + encryptor.finalize()

	if output_file is not None:
		with open(output_file, 'wb') as f:
			f.write(salt + iv + ciphertext)

	return salt + iv + ciphertext, nonce


# Function to decrypt a file
def decrypt_file(input_file, output_file=None):
	with open(input_file, 'rb') as f:
		encrypted_data = f.read()

	salt = encrypted_data[:16]
	iv = encrypted_data[16:32]
	ciphertext = encrypted_data[32:]

	key = derive_key(salt)
	cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=default_backend())
	decryptor = cipher.decryptor()

	plaintext = decryptor.update(ciphertext) + decryptor.finalize()

	if output_file is None:
		return plaintext
	else:
		with open(output_file, 'wb') as f:
			f.write(plaintext)