SSH-key-Manager/add_host.py

# ssh_manager/add_host.py

import os
import subprocess
import glob
from pathlib import Path
from typing import Optional, List, Dict, Set, Tuple
from dataclasses import dataclass
from .utils import print_error, print_warning, print_info, safe_input
from .config import CONF_DIR

@dataclass
class HostConfig:
    label: str
    hostname: str
    user: str = "root"
    port: str = "22"
    identity_file: str = ""

    def to_config_lines(self) -> List[str]:
        """Convert host configuration to SSH config file lines."""
        lines = [
            f"Host {self.label}",
            f"    HostName {self.hostname}",
            f"    User {self.user}",
            f"    Port {self.port}"
        ]
        if self.identity_file:
            lines.append(f"    IdentityFile {self.identity_file}")
        return lines

def get_existing_hosts(conf_dir: str) -> Tuple[Set[str], Dict[str, str]]:
    """
    Get existing host labels and hostnames from config files.
    Returns (host_labels, hostname_to_label) where:
    - host_labels is a set of existing host labels
    - hostname_to_label maps hostnames to their labels
    """
    host_labels = set()
    hostname_to_label = {}
    
    pattern = os.path.join(conf_dir, "*", "config")
    for config_file in glob.glob(pattern):
        try:
            with open(config_file, 'r') as f:
                lines = f.readlines()
                
            current_label = None
            for line in lines:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                    
                if line.lower().startswith('host '):
                    labels = line.split()[1:]
                    for label in labels:
                        if '*' not in label:
                            current_label = label
                            host_labels.add(label)
                            break
                elif current_label and line.lower().startswith('hostname '):
                    hostname = line.split(None, 1)[1].strip()
                    hostname_to_label[hostname] = current_label
                    
        except Exception as e:
            print_warning(f"Error reading config file {config_file}: {e}")
            continue
            
    return host_labels, hostname_to_label

def generate_ssh_key(key_path: Path) -> bool:
    """Generate a new ED25519 SSH key pair."""
    try:
        subprocess.check_call([
            "ssh-keygen",
            "-q",
            "-t", "ed25519",
            "-N", "",
            "-f", str(key_path)
        ])
        print_info(f"Generated new SSH key at {key_path}")
        return True
    except subprocess.CalledProcessError as e:
        print_error(f"Error generating SSH key: {e}")
        return False

def copy_ssh_key(key_path: Path, user: str, hostname: str, port: str) -> bool:
    """Copy SSH public key to remote server."""
    try:
        cmd = ["ssh-copy-id", "-i", str(key_path)]
        if port != "22":
            cmd.extend(["-p", port])
        cmd.append(f"{user}@{hostname}")
        
        subprocess.check_call(cmd)
        print_info("Key successfully copied to remote server.")
        return True
    except subprocess.CalledProcessError as e:
        print_error(f"Error copying key to server: {e}")
        return False

def write_config_file(config_path: Path, config_lines: List[str]) -> bool:
    """Write SSH config lines to file."""
    try:
        config_path.write_text("\n".join(config_lines) + "\n")
        print_info(f"Created/updated config at: {config_path}")
        return True
    except Exception as e:
        print_error(f"Failed to write config to {config_path}: {e}")
        return False

def add_host(conf_dir: str) -> bool:
    """
    Interactive prompt to create a new SSH host in ~/.ssh/conf/<label>/config.
    Offers to generate a new SSH key pair (ed25519) quietly (-q),
    and then prompt to copy that key to the remote server via ssh-copy-id.
    
    Returns True if host was added successfully, False otherwise.
    """
    print_info("Adding a new SSH host...")

    # Get existing hosts to check for duplicates
    existing_labels, hostname_to_label = get_existing_hosts(conf_dir)

    # Get host label
    while True:
        host_label = safe_input("Enter Host label (e.g. myserver): ")
        if not host_label or host_label is None:
            print_error("Host label cannot be empty.")
            return False
            
        host_label = host_label.strip()
        if host_label in existing_labels:
            print_error(f"A host with label '{host_label}' already exists. Please choose a different label.")
            continue
        break
    
    # Get hostname
    while True:
        hostname = safe_input("Enter HostName (IP or domain): ")
        if not hostname or hostname is None:
            print_error("HostName cannot be empty.")
            return False
            
        hostname = hostname.strip()
        if hostname in hostname_to_label:
            existing_label = hostname_to_label[hostname]
            print_error(f"This hostname is already configured for host '{existing_label}'. Please use a different hostname.")
            continue
        break

    # Get optional parameters
    user = safe_input("Enter username (default: 'root'): ") or "root"
    if user is None:
        return False
        
    port = safe_input("Enter SSH port (default: 22): ") or "22"
    if port is None:
        return False

    # Create host configuration
    host_config = HostConfig(
        label=host_label,
        hostname=hostname,
        user=user.strip(),
        port=port.strip()
    )

    # Setup directory structure
    host_dir = Path(conf_dir) / host_label
    if host_dir.exists():
        print_warning(f"Directory {host_dir} already exists; continuing anyway.")
    host_dir.mkdir(mode=0o700, exist_ok=True)
    print_info(f"Created directory: {host_dir}")

    config_path = host_dir / "config"
    if config_path.exists():
        print_warning(f"Config file already exists: {config_path}; it will be overwritten/updated.")

    # Handle SSH key generation
    gen_key_choice = safe_input("Generate a new ed25519 SSH key for this host? (y/n): ")
    if gen_key_choice is None:
        return False

    if gen_key_choice.lower().strip() == 'y':
        key_path = host_dir / "id_ed25519"
        
        if key_path.exists():
            print_warning(f"{key_path} already exists. Skipping generation.")
            host_config.identity_file = str(key_path)
        else:
            if generate_ssh_key(key_path):
                host_config.identity_file = str(key_path)
                
                # Prompt to copy the key
                copy_key = safe_input("Would you like to copy this key to the server now? (y/n): ")
                if copy_key is None:
                    return False
                    
                if copy_key.lower().strip() == 'y':
                    copy_ssh_key(key_path, host_config.user, host_config.hostname, host_config.port)
    else:
        # Handle existing key
        existing_key = safe_input("Enter existing IdentityFile path (or leave empty to skip): ")
        if existing_key is None:
            return False
            
        if existing_key.strip():
            host_config.identity_file = os.path.expanduser(existing_key.strip())

    # Write the config file
    return write_config_file(config_path, host_config.to_config_lines())