import sys, os

from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend


# function to generate a 256-bit symmetric key
def generate_symmetric_key():
	return os.urandom(32)


# function to encrypt data using a symmetric key
def encrypt_symmetric(key, plain_text):
	# generate a random IV
	iv = os.urandom(16)

	# cipher the data using AES in CFB mode
	cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=default_backend())
	encryptor = cipher.encryptor()
	ciphertext = encryptor.update(plain_text) + encryptor.finalize()

	return iv + ciphertext


# function that calls and combines the symmetric and asymmetric encryption
def encrypt_hybrid(public_key, plaintext):
	# generate a random symmetric key
	symmetric_key = generate_symmetric_key()

	encrypted_data = encrypt_symmetric(symmetric_key, plaintext)

	# encrypt the symmetric key with the public key
	encrypted_symmetric_key = public_key.encrypt(
		symmetric_key,
		padding.OAEP(
			mgf=padding.MGF1(algorithm=hashes.SHA256()),
			algorithm=hashes.SHA256(),
			label=None
		)
	)

	# combine the symmetric key and the encrypted data
	return encrypted_symmetric_key + encrypted_data


# main function to encrypt the file
def encrypt_file(public_key, original_file, encrypted_file):
	with open(original_file, 'rb') as f:
		plaintext = f.read()

	encrypted_content = encrypt_hybrid(public_key, plaintext)

	with open(encrypted_file, 'wb') as f:
		f.write(encrypted_content)


# function to load a public key from a file
def load_public_key(file):
	with open(file, 'rb') as key_file:
		public_key = serialization.load_pem_public_key(
			key_file.read(),
		)

	return public_key