112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
# ssh_manager/cli.py
|
|
|
|
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
|
|
from .edit_host import edit_host
|
|
from .list_hosts import list_hosts
|
|
from .regen_key import regenerate_key
|
|
from .remove_host import remove_host
|
|
|
|
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.
|
|
"""
|
|
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 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
|
|
manager.list_hosts()
|
|
|
|
while True:
|
|
manager.display_menu()
|
|
choice = safe_input("Select an option (1-6): ")
|
|
if choice is None:
|
|
continue
|
|
|
|
if not manager.handle_command(choice.strip()):
|
|
break
|
|
|
|
return 0
|