SSH-key-Manager/remove_host.py

# ssh_manager/remove_host.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
"""
Remove host now reuses build_host_list_table to display the same columns:
No. | Host | User | Port | HostName | IP Address | Conf Directory
"""

async def remove_host(conf_dir):
    """
    Remove an SSH host by:
      1) Listing all hosts (same columns as main list)
      2) Letting user pick row number or label
      3) Removing old pub key from remote
      4) Deleting ~/.ssh/conf/<host_label>
    """

    print_info("Remove Host - Step 1: Show current hosts...\n")

    # Reuse the unified table from list_hosts
    headers, final_data = await build_host_list_table(conf_dir)

    if not final_data:
        print_warning("No hosts found. Cannot remove anything.")
        return

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

    # We have final_data rows => need to map row => block
    # So let's gather the raw blocks again to correlate.
    # We'll do a separate approach or we can parse final_data.
    # Easiest: Re-run load_config_file if needed or:
    blocks = []
    # The last gather call for build_host_list_table used load_config_file already
    # but it doesn't return the correlation. We'll replicate the logic quickly.

    pattern = os.path.join(conf_dir, "*", "config")
    conf_files = sorted(os.listdir(conf_dir))

    # Actually, let's do a small approach:
    from .list_hosts import gather_host_info, sort_by_ip
    all_blocks = []
    import glob

    for cfile in glob.glob(os.path.join(conf_dir, "*", "config")):
        blocks.extend(load_config_file(cfile))

    # gather the same big list
    results = await gather_host_info(blocks)
    sorted_rows = sort_by_ip(results)
    # sorted_rows is a list of 7-tuples:
    # (host_label, user, colored_port, hostname, colored_ip, conf_display, raw_ip)

    choice = safe_input("Enter the row number or Host label to remove: ")
    if choice is None:
        return
    choice = choice.strip()
    if not choice:
        print_error("Invalid empty choice.")
        return

    target_tuple = None

    # If digit => index in sorted_rows
    if choice.isdigit():
        idx = int(choice)
        if idx < 1 or idx > len(sorted_rows):
            print_warning(f"Row number {idx} is invalid.")
            return
        target_tuple = sorted_rows[idx - 1]
    else:
        # They typed a label
        for t in sorted_rows:
            if t[0] == choice:  # t[0] = host_label
                target_tuple = t
                break
        if not target_tuple:
            print_warning(f"No matching host label '{choice}' found.")
            return

    # target_tuple is (host_label, user, colored_port, hostname, colored_ip, conf_dir, raw_ip)
    host_label = target_tuple[0]
    hostname   = target_tuple[3]
    user       = target_tuple[1]
    # parse port from colored_port? We can do a quick approach. Alternatively, we re-run the load_config approach again.
    # We do a second approach: let's see if we can find the block in "blocks".
    # But let's parse the config block to get the real port, identity, etc.

    # We find the matching block in "blocks" by host_label
    found_block = None
    for b in 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", "")

    # now do removal logic
    if not host_label:
        print_warning("Target block has no Host label. Cannot remove.")
        return

    if identity_file and identity_file != "N/A":
        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}")

        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)

    # remove local folder
    host_dir = os.path.join(conf_dir, host_label)
    import shutil
    if os.path.isdir(host_dir):
        confirm = safe_input(f"Are you sure you want to delete local folder {host_dir}? (y/n): ")
        if confirm and confirm.lower().startswith('y'):
            try:
                shutil.rmtree(host_dir)
                print_info(f"Removed local folder: {host_dir}")
            except Exception as e:
                print_warning(f"Could not remove folder {host_dir}: {e}")
        else:
            print_warning("Local folder not removed.")
    else:
        print_warning(f"Local folder {host_dir} not found. Nothing to remove.")

async def remove_old_key_remote(old_pub_data, user, hostname, port):
    # same logic as before
    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).")