import os
import base64
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 import padding
from cryptography.hazmat.primitives.asymmetric import padding as asymetric_padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

class SymetricCryptomanager:
    def __init__(self, path_to_key):
        if os.path.exists(path_to_key):
            with open(path_to_key, "rb") as file:
                self.key = file.read()
        else:
            with open(path_to_key, "wb") as file:
                self.key = os.urandom(32)
                file.write(self.key)

    def encrypt_message(self, message):
        # Pad the message to make its size a multiple of block size
        padder = padding.PKCS7(algorithms.AES.block_size).padder()  # AES block size is 128 bits
        padded_message = padder.update(message) + padder.finalize()

        iv = os.urandom(16)  # Initialization vector
        cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv))
        encryptor = cipher.encryptor()
        tmp = encryptor.update(padded_message) + encryptor.finalize()
        return base64.b64encode(iv + tmp)

    def decrypt_message(self, encoded_message):
        try:
            decoded_data = base64.b64decode(encoded_message)  # Decode the base64 data
            iv = decoded_data[:16]
            cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv))
            decryptor = cipher.decryptor()
            padded_plaintext = decryptor.update(decoded_data[16:]) + decryptor.finalize()

            # Unpad the plaintext
            unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
            plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()
            return plaintext.decode('utf-8')
        except Exception as e:
            print(f"Decryption error: {e}")
            return ""


class AsymetricCryptoManager:

    def __init__(self, name, path_to_private_key):
        if os.path.exists(path_to_private_key):
            with open(path_to_private_key, "rb") as key_file:
                self.private_key = serialization.load_pem_private_key(
                    key_file.read(),
                    password=None,
                )
        else:
            self.generate_and_write_keys(path_to_private_key, '/'.join(path_to_private_key.split('/')[:-1]) + f'/{name}.pem')


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

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

        public_key = self.private_key.public_key().public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )

        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 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(),
                backend=default_backend()
            )

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

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

            return text

        except:
            return ""

# crypto = AsymetricCryptoManager('denis', '/home/ubuntu/abradabra/private_1.pem')
# msg = 'hello'
# print(crypto.decrypt_message(crypto.encrypt_message_with_key_from_file('denis.pem', msg.encode())))

# crypto = SymetricCryptomanager('sym.key')
# msg = 'hello sss'
# print(crypto.decrypt_message(crypto.encrypt_message(msg.encode())))