refactor

cli.py

@@ -2,6 +2,8 @@
 
 import os
 import asyncio
+from typing import Callable, Dict, Optional
+from functools import wraps
 from .utils import print_info, print_error, print_warning, Colors, safe_input
 from .config import SSH_DIR, CONF_DIR, SOCKET_DIR, MAIN_CONFIG, DEFAULT_CONFIG_CONTENT
 from .add_host import add_host
@@ -10,62 +12,101 @@ from .list_hosts import list_hosts
 from .regen_key import regenerate_key
 from .remove_host import remove_host
 
-def ensure_ssh_setup():
+def ensure_ssh_setup() -> None:
     """
     Creates ~/.ssh, ~/.ssh/conf, and ~/.ssh/s if missing,
     and writes a default ~/.ssh/config if it doesn't exist.
     """
-    if not os.path.isdir(SSH_DIR):
-        os.makedirs(SSH_DIR, mode=0o700, exist_ok=True)
-        print_info(f"Created directory: {SSH_DIR}")
-
-    if not os.path.isdir(CONF_DIR):
-        os.makedirs(CONF_DIR, mode=0o700, exist_ok=True)
-        print_info(f"Created directory: {CONF_DIR}")
-
-    if not os.path.isdir(SOCKET_DIR):
-        os.makedirs(SOCKET_DIR, mode=0o700, exist_ok=True)
-        print_info(f"Created directory: {SOCKET_DIR}")
+    directories = [
+        (SSH_DIR, "SSH directory"),
+        (CONF_DIR, "configuration directory"),
+        (SOCKET_DIR, "socket directory")
+    ]
+    
+    for directory, description in directories:
+        if not os.path.isdir(directory):
+            os.makedirs(directory, mode=0o700, exist_ok=True)
+            print_info(f"Created {description}: {directory}")
 
     if not os.path.isfile(MAIN_CONFIG):
         with open(MAIN_CONFIG, "w") as f:
             f.write(DEFAULT_CONFIG_CONTENT)
         print_info(f"Created default SSH config at: {MAIN_CONFIG}")
 
-def main():
-    ensure_ssh_setup()
+def async_handler(func: Callable) -> Callable:
+    """Decorator to handle async functions in the command dispatch"""
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        return asyncio.run(func(*args, **kwargs))
+    return wrapper
 
+class SSHManager:
+    def __init__(self):
+        self.commands: Dict[str, tuple[Callable, str]] = {
+            "1": (self.list_hosts, "List Hosts"),
+            "2": (self.add_host, "Add a Host"),
+            "3": (self.edit_host, "Edit a Host"),
+            "4": (self.regenerate_key, "Regenerate Key"),
+            "5": (self.remove_host, "Remove Host"),
+            "6": (self.exit_app, "Exit")
+        }
+
+    @async_handler
+    async def list_hosts(self) -> None:
+        await list_hosts(CONF_DIR)
+
+    def add_host(self) -> None:
+        add_host(CONF_DIR)
+
+    @async_handler
+    async def edit_host(self) -> None:
+        await edit_host(CONF_DIR)
+
+    @async_handler
+    async def regenerate_key(self) -> None:
+        await regenerate_key(CONF_DIR)
+
+    @async_handler
+    async def remove_host(self) -> None:
+        await remove_host(CONF_DIR)
+
+    def exit_app(self) -> None:
+        print_info("Exiting...")
+        raise SystemExit(0)
+
+    def display_menu(self) -> None:
+        print("\n" + f"{Colors.CYAN}{Colors.BOLD}SSH Config Manager Menu{Colors.RESET}")
+        for key, (_, description) in self.commands.items():
+            print(f"{key}. {description}")
+
+    def handle_command(self, choice: str) -> bool:
+        if choice not in self.commands:
+            print_error("Invalid choice. Please select 1 through 6.")
+            return True
+        
+        try:
+            self.commands[choice][0]()
+            return choice != "6"
+        except SystemExit:
+            return False
+        except Exception as e:
+            print_error(f"Error executing command: {str(e)}")
+            return True
+
+def main() -> int:
+    ensure_ssh_setup()
+    manager = SSHManager()
+    
     # Display the server list on first load
-    asyncio.run(list_hosts(CONF_DIR))
+    manager.list_hosts()
 
     while True:
-        print("\n" + f"{Colors.CYAN}{Colors.BOLD}SSH Config Manager Menu{Colors.RESET}")
-        print("1. List Hosts")
-        print("2. Add a Host")
-        print("3. Edit a Host")
-        print("4. Regenerate Key")
-        print("5. Remove Host")
-        print("6. Exit")
-
+        manager.display_menu()
         choice = safe_input("Select an option (1-6): ")
         if choice is None:
-            continue  # User pressed Ctrl+C => safe_input returns None => re-show menu
-
-        choice = choice.strip()
-        if choice == '1':
-            asyncio.run(list_hosts(CONF_DIR))
-        elif choice == '2':
-            add_host(CONF_DIR)
-        elif choice == '3':
-            asyncio.run(edit_host(CONF_DIR))
-        elif choice == '4':
-            asyncio.run(regenerate_key(CONF_DIR))
-        elif choice == '5':
-            asyncio.run(remove_host(CONF_DIR))
-        elif choice == '6':
-            print_info("Exiting...")
+            continue
+        
+        if not manager.handle_command(choice.strip()):
             break
-        else:
-            print_error("Invalid choice. Please select 1 through 6.")
 
     return 0