diff --git a/main.py b/main.py index 6b31855..f43d8b5 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,11 @@ import base64 +import os from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization, hashes from cryptography.hazmat.primitives.asymmetric import rsa, padding -# Функции для генерации ключей, шифрования и расшифрования +PRIVATE_KEY_FILE = 'private_key.pem' +PUBLIC_KEY_FILE = 'public_key.pem' def generate_keys(): private_key = rsa.generate_private_key( @@ -14,6 +16,44 @@ def generate_keys(): public_key = private_key.public_key() return private_key, public_key +def save_keys(private_key, public_key): + # Сохранение закрытого ключа + with open(PRIVATE_KEY_FILE, 'wb') as f: + f.write(private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption() + )) + + # Сохранение открытого ключа + with open(PUBLIC_KEY_FILE, 'wb') as f: + f.write(public_key.public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo + )) + +def load_keys(): + # Загрузка закрытого ключа + with open(PRIVATE_KEY_FILE, 'rb') as f: + private_key_data = f.read() + + private_key = serialization.load_pem_private_key( + private_key_data, + password=None, + backend=default_backend() + ) + + # Загрузка открытого ключа + with open(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 private_key, public_key + def encrypt_message(message, public_key): encrypted = public_key.encrypt( message.encode(), @@ -23,11 +63,9 @@ def encrypt_message(message, public_key): label=None ) ) - # Кодируем в Base64 для представления в виде строки return base64.b64encode(encrypted).decode('utf-8') def decrypt_message(encrypted_message, private_key): - # Декодируем из Base64 обратно в байты encrypted_bytes = base64.b64decode(encrypted_message) decrypted = private_key.decrypt( encrypted_bytes, @@ -39,8 +77,22 @@ def decrypt_message(encrypted_message, private_key): ) return decrypted.decode() -# Генерация ключей -private_key, public_key = generate_keys() +# Запрос выбора действия +action = input("Хотите сгенерировать новые ключи (1) или использовать существующие (2)? ") + +if action == '1': + # Генерация новых ключей + private_key, public_key = generate_keys() + save_keys(private_key, public_key) + print(f"Новые ключи сгенерированы и сохранены в '{PRIVATE_KEY_FILE}' и '{PUBLIC_KEY_FILE}'.") +elif action == '2': + # Загрузка существующих ключей + if os.path.exists(PRIVATE_KEY_FILE) and os.path.exists(PUBLIC_KEY_FILE): + private_key, public_key = load_keys() + print(f"Существующие ключи загружены из '{PRIVATE_KEY_FILE}' и '{PUBLIC_KEY_FILE}'.") + else: + print(f"Файлы ключей не найдены. Пожалуйста, сгенерируйте новые ключи.") + exit() # Шифрование сообщения message = "Это лицензионный ключ."