SSH-key-Manager/regen_key.py

import os
import subprocess
import asyncio
from .utils import (
    print_info,
    print_warning,
    print_error,
    safe_input
)
from .list_hosts import (
    build_host_list_table,
    load_config_file,
    gather_host_info,
    sort_by_ip
)

async def regenerate_key(conf_dir):
    """
    Regenerate the SSH key for a chosen host by:
      1) Displaying the unified table of hosts (No. | Host | User | Port | HostName | IP Address | Conf Directory).
      2) Letting you pick a row number or host label.
      3) Reading/deleting any existing local keys,
      4) Generating a new key,
      5) Optionally copying it to the remote,
      6) Removing the old pub key from the remote authorized_keys if present.
    """

    print_info("Regenerate Key - Step 1: Show current hosts...\n")

    # 1) Reuse the same columns as the main 'list_hosts'
    headers, final_data = await build_host_list_table(conf_dir)
    if not final_data:
        print_warning("No hosts found. Cannot regenerate a key.")
        return

    from tabulate import tabulate
    print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
    print(tabulate(final_data, headers=headers, tablefmt="grid"))

    # 2) We need to correlate row => actual config block.
    #    We'll replicate the logic that build_host_list_table uses.
    all_blocks = []
    import glob
    for cfile in glob.glob(os.path.join(conf_dir, "*", "config")):
        all_blocks.extend(load_config_file(cfile))

    results = await gather_host_info(all_blocks)
    sorted_rows = sort_by_ip(results)
    # sorted_rows is a list of 7-tuples:
    # (host_label, user, colored_port, hostname, colored_ip, conf_path, raw_ip)

    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_tuple = None
    if choice.isdigit():
        row_idx = int(choice)
        if row_idx < 1 or row_idx > len(sorted_rows):
            print_warning(f"Invalid row number {row_idx}.")
            return
        target_tuple = sorted_rows[row_idx - 1]
    else:
        # user typed a label
        for t in sorted_rows:
            if t[0] == choice:  # t[0] is host_label
                target_tuple = t
                break
        if not target_tuple:
            print_warning(f"No matching host label '{choice}' found.")
            return

    # 3) Retrieve the full config block so we have Port, IdentityFile, etc.
    host_label = target_tuple[0]
    hostname   = target_tuple[3]  # (host_label, user, colored_port, hostname, ...)
    user       = target_tuple[1]

    # find that block
    found_block = None
    for b in all_blocks:
        if b.get("Host") == host_label:
            found_block = b
            break

    if not found_block:
        print_warning(f"No config block found for '{host_label}'.")
        return

    port_str      = found_block.get("Port", "22")
    port          = int(port_str)
    identity_file = found_block.get("IdentityFile", "")

    # If missing or "N/A", we can't regenerate
    if not identity_file or identity_file == "N/A":
        print_error("No IdentityFile found in config; cannot regenerate.")
        return

    # 4) Remove old local key files
    expanded_key  = os.path.expanduser(identity_file)
    pub_path      = expanded_key + ".pub"
    old_pub_data  = ""

    if os.path.isfile(pub_path):
        try:
            with open(pub_path, "r") as f:
                old_pub_data = f.read().rstrip("\n")
        except Exception as e:
            print_warning(f"Could not read old pub key: {e}")

    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
    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 the new key to remote if user wants
    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 we had old_pub_data
    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):
    """
    Checks and removes the exact matching line from remote authorized_keys
    via grep -Fxq / grep -vFx.
    """
    check_cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no",
        "-p", str(port),
        f"{user}@{hostname}",
        f'grep -Fxq "{old_pub_data}" ~/.ssh/authorized_keys'
    ]
    found_key = False
    try:
        subprocess.check_call(check_cmd)
        found_key = True
    except subprocess.CalledProcessError:
        pass

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

    remove_cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no",
        "-p", str(port),
        f"{user}@{hostname}",
        f'grep -vFx "{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 (permissions or other error).")