Рефакторинг кода

.gitignore

@@ -1,2 +1,3 @@
 /.venv/
 /keys/
+*.pem

main.py

@@ -1,7 +1,7 @@
 import os
-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
+from utilits.encrypt import encrypt_message
+from utilits.decrypt import decrypt_message
+from utilits.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'
@@ -10,15 +10,13 @@ PUBLIC_KEY_FILE = f'{KEYS_DIR}/public-key.pem'
 if not os.path.exists(KEYS_DIR):
     os.makedirs(KEYS_DIR)
 
-action = input("Хотите сгенерировать новые ключи (1) или использовать существующие (2)? ")
-
-if action == '1':
+if not os.path.exists(PRIVATE_KEY_FILE) and not os.path.exists(PUBLIC_KEY_FILE):
     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':
+else:
     if os.path.exists(PRIVATE_KEY_FILE) and os.path.exists(PUBLIC_KEY_FILE):
         pass
     else:

utilits/decrypt.py

@@ -3,6 +3,7 @@ from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.asymmetric import padding
 
+
 def load_keys_private(private_key_file):
     with open(private_key_file, 'rb') as f:
         private_key_data = f.read()
@@ -14,6 +15,7 @@ def load_keys_private(private_key_file):
     )
     return private_key
 
+
 def decrypt_message(encrypted_message, private_key_file):
     private_key = load_keys_private(private_key_file)
     encrypted_bytes = base64.b64decode(encrypted_message)

utilits/encrypt.py

@@ -26,4 +26,3 @@ def encrypt_message(message, public_key_file):
         )
     )
     return base64.b64encode(encrypted).decode('utf-8')
-

utilits/generate_save.py

@@ -32,12 +32,3 @@ def save_public_key(public_key, public_key_file):
             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')