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

# Создать пару ключей
def generate_key_pair():
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    public_key = private_key.public_key()
    return private_key, public_key

# Хранить закрытый ключ
def save_private_key(private_key, filename="private_key.pem"):
    with open(filename, "wb") as key_file:
        key_file.write(private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption()
        ))
        
# Загрузить закрытый ключ
def load_private_key(filename="private_key.pem"):
    with open(filename, "rb") as key_file:
        private_key = serialization.load_pem_private_key(
            key_file.read(),
            password=None,
            backend=default_backend()
        )
    return private_key

# Хранить открытый ключ
def save_public_key(public_key, filename="public_key.pem"):
    with open(filename, "wb") as key_file:
        key_file.write(public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ))

#Загрузить открытый ключ
def load_public_key(filename="public_key.pem"):
    with open(filename, "rb") as key_file:
        public_key = serialization.load_pem_public_key(
            key_file.read(),
            backend=default_backend()
        )
    return public_key
    
# Подпись закрытым ключом
def sign_data(private_key, data):
    signature = private_key.sign(
        data,
        padding.PKCS1v15(),
        hashes.SHA256()
    )
    return signature

# Проверка подписи открытым ключом
def verify_signature(public_key, signature, data):
    try:
        public_key.verify(
            signature,
            data,
            padding.PKCS1v15(),
            hashes.SHA256()
        )
        return True
    except InvalidSignature:
        return False

# Пример использования
private_key, public_key = generate_key_pair()

# Сохранить пару ключей
save_private_key(private_key)
save_public_key(public_key)

# Аффидевит
original_data = b"Hello, World!"
signature = sign_data(private_key, original_data)

# Передать оригинальные данные и подпись
# ... (В реальных приложениях эти два значения могут передаваться куда угодно)

# Проверка
loaded_public_key = load_public_key()
is_valid_signature = verify_signature(loaded_public_key, signature, original_data)

if is_valid_signature:
    print("Проверка подписи прошла успешно, данные заполнены.")
else:
    print("Проверка подписи не удалась, возможно, данные были подделаны.")