Разграниченеи загрузки ключей

main.py

@@ -1,11 +1,16 @@
-import base64
 import os
+import base64
 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'
+KEYS_DIR = 'keys'
-PUBLIC_KEY_FILE = 'public_key.pem'
+PRIVATE_KEY_FILE = f'{KEYS_DIR}/private-key.pem'
+PUBLIC_KEY_FILE = f'{KEYS_DIR}/public-key.pem'
+
+if not os.path.exists(KEYS_DIR):
+    os.makedirs(KEYS_DIR)
+
 
 def generate_keys():
     private_key = rsa.generate_private_key(
@@ -16,6 +21,7 @@ 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:
@@ -32,7 +38,8 @@ def save_keys(private_key, public_key):
             format=serialization.PublicFormat.SubjectPublicKeyInfo
         ))
 
-def load_keys():
+
+def load_keys_private():
     # Загрузка закрытого ключа
     with open(PRIVATE_KEY_FILE, 'rb') as f:
         private_key_data = f.read()
@@ -43,6 +50,10 @@ def load_keys():
         backend=default_backend()
     )
 
+    return private_key
+
+
+def load_keys_public():
     # Загрузка открытого ключа
     with open(PUBLIC_KEY_FILE, 'rb') as f:
         public_key_data = f.read()
@@ -52,9 +63,11 @@ def load_keys():
         backend=default_backend()
     )
 
-    return private_key, public_key
+    return public_key
 
-def encrypt_message(message, public_key):
+
+def encrypt_message(message):
+    public_key = load_keys_public()
     encrypted = public_key.encrypt(
         message.encode(),
         padding.OAEP(
@@ -65,7 +78,9 @@ def encrypt_message(message, public_key):
     )
     return base64.b64encode(encrypted).decode('utf-8')
 
-def decrypt_message(encrypted_message, private_key):
+
+def decrypt_message(encrypted_message):
+    private_key = load_keys_private()
     encrypted_bytes = base64.b64decode(encrypted_message)
     decrypted = private_key.decrypt(
         encrypted_bytes,
@@ -77,6 +92,7 @@ def decrypt_message(encrypted_message, private_key):
     )
     return decrypted.decode()
 
+
 # Запрос выбора действия
 action = input("Хотите сгенерировать новые ключи (1) или использовать существующие (2)? ")
 
@@ -88,17 +104,16 @@ if action == '1':
 elif action == '2':
     # Загрузка существующих ключей
     if os.path.exists(PRIVATE_KEY_FILE) and os.path.exists(PUBLIC_KEY_FILE):
-        private_key, public_key = load_keys()
+        pass
-        print(f"Существующие ключи загружены из '{PRIVATE_KEY_FILE}' и '{PUBLIC_KEY_FILE}'.")
     else:
         print(f"Файлы ключей не найдены. Пожалуйста, сгенерируйте новые ключи.")
         exit()
 
 # Шифрование сообщения
 message = "Это лицензионный ключ."
-encrypted_message = encrypt_message(message, public_key)
+encrypted_message = encrypt_message(message)
 print("Зашифрованное сообщение (Base64):", encrypted_message)
 
 # Дешифрование сообщения
-decrypted_message = decrypt_message(encrypted_message, private_key)
+decrypted_message = decrypt_message(encrypted_message)
 print("Расшифрованное сообщение:", decrypted_message)