Генерация вынесена в отдельные функции

generate_save.py

@@ -0,0 +1,43 @@
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+
+def generate_private_key():
+    private_key = rsa.generate_private_key(
+        public_exponent=65537,
+        key_size=2048,
+        backend=default_backend()
+    )
+    return private_key
+
+
+def generate_public_key(private_key):
+    public_key = private_key.public_key()
+    return public_key
+
+
+def save_private_key(private_key, private_key_file):
+    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()
+        ))
+
+
+def save_public_key(public_key, public_key_file):
+    with open(public_key_file, 'wb') as f:
+        f.write(public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        ))
+
+
+# Пример использования
+if __name__ == "__main__":
+    private_key = generate_private_key()
+    public_key = generate_public_key(private_key)
+
+    save_private_key(private_key, 'private_key.pem')
+    save_public_key(public_key, 'public_key.pem')

main.py

@@ -1,9 +1,7 @@
 import os
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric import rsa
 from encrypt import encrypt_message
 from decrypt import decrypt_message
+from generate_save import generate_private_key, generate_public_key, save_private_key, save_public_key
 
 KEYS_DIR = 'keys'
 PRIVATE_KEY_FILE = f'{KEYS_DIR}/private-key.pem'
@@ -12,44 +10,15 @@ 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(
-        public_exponent=65537,
-        key_size=2048,
-        backend=default_backend()
-    )
-    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
-        ))
-
-
-# Запрос выбора действия
 action = input("Хотите сгенерировать новые ключи (1) или использовать существующие (2)? ")
 
 if action == '1':
-    # Генерация новых ключей
-    private_key, public_key = generate_keys()
-    save_keys(private_key, public_key)
+    private_key = generate_private_key()
+    save_private_key(private_key, PRIVATE_KEY_FILE)
+    public_key = generate_public_key(private_key)
+    save_public_key(public_key, PUBLIC_KEY_FILE)
     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):
         pass
     else: