import base64
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes

class Encryptor:
    """
    Class to encrypt/decrypt content
    """
    def __init__(self, password: str):
        self.password = password.encode()
        self.salt_size = 16
        self.iv_size = 16
        self.key_size = 32
        self.iterations = 100_000

    def _derive_key(self, salt: bytes) -> bytes:
        return PBKDF2(self.password, salt, dkLen=self.key_size, count=self.iterations)

    def encrypt(self, plaintext: str | bytes) -> str:
        """
        Encrypte une chaîne de texte en base64.
        """
        if isinstance(plaintext, str):
            plaintext_bytes = plaintext.encode()
        else:
            plaintext_bytes = plaintext

        salt = get_random_bytes(self.salt_size)
        iv = get_random_bytes(self.iv_size)
        key = self._derive_key(salt)

        # Padding (PKCS7)
        pad_len = AES.block_size - (len(plaintext_bytes) % AES.block_size)
        padded = plaintext_bytes + bytes([pad_len] * pad_len)

        cipher = AES.new(key, AES.MODE_CBC, iv)
        ciphertext = cipher.encrypt(padded)

        encrypted_data = base64.b64encode(salt + iv + ciphertext)
        return encrypted_data

    def decrypt(self, encrypted_text: str) -> bytes:
        """
        Décrypte une chaîne encodée en base64.
        """
        data = base64.b64decode(encrypted_text)
        salt = data[:self.salt_size]
        iv = data[self.salt_size:self.salt_size + self.iv_size]
        ciphertext = data[self.salt_size + self.iv_size:]

        key = self._derive_key(salt)
        cipher = AES.new(key, AES.MODE_CBC, iv)
        padded_plaintext = cipher.decrypt(ciphertext)

        # Remove padding
        pad_len = padded_plaintext[-1]
        plaintext = padded_plaintext[:-pad_len]
        return plaintext  # retourne des bytes