Добавлено создание ключей

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 = "Это лицензионный ключ."