"""
Encrypt messages encrypted with RSA.
"""

import base64
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding


class RSAEncryptor:
    """
    Module for encrypting messages encrypted with RSA.

    Provides
    --------

    *   :class:`RSAEncryptor` - class for encrypting messages.
    """
    def __init__(self, public_key_file):
        self.public_key_file = public_key_file
        self.public_key = self.load_public_key()

    def load_public_key(self):
        """
        Loads a public key from a file.

        Parameters
        ----------
        self : `RSAEncryptor`
            The object itself.

        Returns
        -------
        `cryptography.hazmat.primitives.asymmetric.rsa.RSAPublicKey`
            The loaded key.
        """
        with open(self.public_key_file, 'rb') as f:
            public_key_data = f.read()

        public_key = serialization.load_pem_public_key(
            public_key_data,
            backend=default_backend()
        )
        return public_key

    def encrypt_message(self, message):
        """
        Encrypts a message using RSA public key encryption.

        Parameters
        ----------
        message : str
            The message to be encrypted.

        Returns
        -------
        str
            The encrypted message encoded in Base64.
        """
        encrypted = self.public_key.encrypt(
            message.encode(),
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return base64.b64encode(encrypted).decode('utf-8')