SSH-key-Manager/regen_key.py

import os
import glob
import subprocess
import asyncio
from collections import OrderedDict
from .utils import (
    print_info,
    print_warning,
    print_error,
    safe_input,
    Colors
)
from .list_hosts import load_config_file, check_ssh_port

async def get_all_host_blocks(conf_dir):
    pattern    = os.path.join(conf_dir, "*", "config")
    conf_files = sorted(glob.glob(pattern))

    all_blocks = []
    for conf_file in conf_files:
        blocks = load_config_file(conf_file)
        all_blocks.extend(blocks)

    if not all_blocks:
        return []
    return all_blocks

async def regenerate_key(conf_dir):
    """
    Menu-driven function to regenerate a key for a selected host:
      1) Show the host table with row numbers
      2) Let user pick row # or label
      3) Read old pub data
      4) Remove old local key files
      5) Generate new key
      6) Copy new key
      7) Remove old key from remote authorized_keys (improved logic to detect existence)
    """
    print_info("Regenerate Key - Step 1: Show current hosts...\n")

    # 1) Gather host blocks
    all_blocks = await get_all_host_blocks(conf_dir)
    if not all_blocks:
        print_warning("No hosts found. Cannot regenerate a key.")
        return

    # Display them in a table (similar to list_hosts):
    import socket
    from tabulate import tabulate

    table_rows = []
    for idx, block in enumerate(all_blocks, start=1):
        host_label    = block.get("Host", "N/A")
        hostname      = block.get("HostName", "N/A")
        user          = block.get("User", "N/A")
        port          = int(block.get("Port", "22"))
        identity_file = block.get("IdentityFile", "N/A")

        # Check IP
        try:
            ip_address = socket.gethostbyname(hostname)
            port_open  = await asyncio.wait_for(check_ssh_port(ip_address, port), timeout=1)
        except:
            ip_address = None
            port_open  = False

        ip_disp   = f"\033[0;32m{ip_address}\033[0m" if ip_address else "\033[0;31mN/A\033[0m"
        port_disp = f"\033[0;32m{port}\033[0m" if port_open else f"\033[0;31m{port}\033[0m"

        row = [
            idx,
            host_label,
            user,
            port_disp,
            hostname,
            ip_disp,
            identity_file
        ]
        table_rows.append(row)

    headers = ["No.", "Host", "User", "Port", "HostName", "IP", "IdentityFile"]
    print("\nSSH Conf Subdirectory Host List")
    print(tabulate(table_rows, headers=headers, tablefmt="grid"))

    # 2) Prompt for row # or label
    choice = safe_input("Enter the row number or the Host label to regenerate: ")
    if choice is None:
        return
    choice = choice.strip()
    if not choice:
        print_error("No choice given.")
        return

    target_block = None
    if choice.isdigit():
        row_idx = int(choice)
        if row_idx < 1 or row_idx > len(all_blocks):
            print_warning(f"Invalid row number {row_idx}.")
            return
        target_block = all_blocks[row_idx - 1]
    else:
        # user typed a label
        for b in all_blocks:
            if b.get("Host") == choice:
                target_block = b
                break
        if not target_block:
            print_warning(f"No matching host label '{choice}' found.")
            return

    # 3) Gather info from block
    host_label    = target_block.get("Host", "")
    hostname      = target_block.get("HostName", "")
    user          = target_block.get("User", "root")
    port          = int(target_block.get("Port", "22"))
    identity_file = target_block.get("IdentityFile", "")

    if not host_label or not hostname:
        print_error("Missing Host or HostName; cannot regenerate.")
        return
    if not identity_file or identity_file == "N/A":
        print_error("No IdentityFile found in config; cannot regenerate.")
        return

    # Derive local paths
    expanded_key  = os.path.expanduser(identity_file)
    key_dir       = os.path.dirname(expanded_key)
    pub_path      = expanded_key + ".pub"

    # 3a) Read old pub key data
    old_pub_data = ""
    if os.path.isfile(pub_path):
        try:
            with open(pub_path, "r") as f:
                old_pub_data = f.read().strip()
        except Exception as e:
            print_warning(f"Could not read old pub key: {e}")
    else:
        print_warning("No old pub key found locally.")

    # 4) Remove old local key files
    print_info("Removing old key files locally...")
    for path in [expanded_key, pub_path]:
        if os.path.isfile(path):
            try:
                os.remove(path)
                print_info(f"Removed {path}")
            except Exception as e:
                print_warning(f"Could not remove {path}: {e}")

    # 5) Generate new key
    print_info("Generating new ed25519 SSH key...")
    new_key_path = expanded_key  # Reuse the same path from config
    cmd = ["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-f", new_key_path]
    try:
        subprocess.check_call(cmd)
        print_info(f"Generated new SSH key at {new_key_path}")
    except subprocess.CalledProcessError as e:
        print_error(f"Error generating new SSH key: {e}")
        return

    # 6) Copy new key to remote
    copy_choice = safe_input("Copy new key to remote now? (y/n): ")
    if copy_choice and copy_choice.lower().startswith('y'):
        ssh_copy_cmd = ["ssh-copy-id", "-i", new_key_path]
        if port != 22:
            ssh_copy_cmd += ["-p", str(port)]
        ssh_copy_cmd.append(f"{user}@{hostname}")
        try:
            subprocess.check_call(ssh_copy_cmd)
            print_info("New key successfully copied to remote server.")
        except subprocess.CalledProcessError as e:
            print_error(f"Error copying new key: {e}")

    # 7) Remove old key from authorized_keys if old_pub_data is non-empty
    if old_pub_data:
        print_info("Attempting to remove old key from remote authorized_keys...")
        await remove_old_key_remote(old_pub_data, user, hostname, port)
    else:
        print_warning("No old pub key data found locally, skipping remote removal.")

async def remove_old_key_remote(old_pub_data, user, hostname, port):
    """
    1) Check if old_pub_data is present on remote server (grep -q).
    2) If found, remove it with grep -v ...
    3) Otherwise, print "No old key found on remote."
    """
    # 1) Check if old_pub_data exists in authorized_keys
    check_cmd = [
        "ssh",
        "-o", "StrictHostKeyChecking=no",
        "-p", str(port),
        f"{user}@{hostname}",
        f"grep -F '{old_pub_data}' ~/.ssh/authorized_keys"
    ]
    found_key = False

    try:
        subprocess.check_call(check_cmd)
        found_key = True
    except subprocess.CalledProcessError:
        # grep returns exit code 1 if not found
        pass

    if not found_key:
        print_warning("No old key found on remote authorized_keys.")
        return

    # 2) Actually remove it
    remove_cmd = [
        "ssh",
        "-o", "StrictHostKeyChecking=no",
        "-p", str(port),
        f"{user}@{hostname}",
        f"grep -v '{old_pub_data}' ~/.ssh/authorized_keys > ~/.ssh/tmp && mv ~/.ssh/tmp ~/.ssh/authorized_keys"
    ]
    try:
        subprocess.check_call(remove_cmd)
        print_info("Old public key removed from remote authorized_keys.")
    except subprocess.CalledProcessError:
        print_warning("Failed to remove old key from remote authorized_keys (permission or other error).")