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

class CryptoManager:
    def __init__(self):
        self.private_key = ""
        self.public_key = ""

    def generate_and_write_keys(self, path_to_private_key, path_to_public_key):
        key = rsa.generate_private_key(
            backend=default_backend(),
            public_exponent=65537,
            key_size=2048
        )

        private_key = key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption(),
        )

        public_key = key.public_key().public_bytes(
            serialization.Encoding.OpenSSH,
            serialization.PublicFormat.OpenSSH
        )

        with open(path_to_private_key, 'w') as file:
            file.write(private_key.decode('utf-8'))

        with open(path_to_public_key, 'w') as file:
            file.write(public_key.decode('utf-8'))


    def load_keys(self, path_to_key_file):
        with open(path_to_key_file, "rb") as key_file:
            key = serialization.load_pem_private_key(
                key_file.read(),
                password=None,
            )

            self.private_key = key.private_bytes(
                serialization.Encoding.PEM,
                serialization.PrivateFormat.PKCS8,
                serialization.NoEncryption(),
            )

            self.public_key = key.public_key().public_bytes(
                serialization.Encoding.OpenSSH,
                serialization.PublicFormat.OpenSSH
            )

            self.key = key
            self.pubkey = key.public_key()

    def encrypt_message_with_key_from_file(self, path_to_public_key, message):
        with open(path_to_public_key, "rb") as key_file:
            pubkey = serialization.load_pem_public_key(
                key_file.read(),
                password=None,
            )

        return pubkey.encrypt(
            message,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )

    def decrypt_message(self, encoded_message): # обработка когда не дешифруется
        text = self.key.decrypt(
            encoded_message,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        ).decode('utf-8')

        return text