SSH-key-Manager/remove_host.py

# ssh_manager/remove_host.py

import os
import glob
import subprocess
import asyncio
from collections import OrderedDict
from .utils import (
    print_info,
    print_warning,
    print_error,
    safe_input
)
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)

    return all_blocks

async def remove_host(conf_dir):
    """
    Remove an SSH host by:
      1) Listing all hosts
      2) Letting user pick row number or label
      3) Attempting to remove the old pub key from remote authorized_keys
      4) Deleting the subdirectory in ~/.ssh/conf/<host_label>
    """
    print_info("Remove Host - Step 1: Show current hosts...\n")

    all_blocks = await get_all_host_blocks(conf_dir)
    if not all_blocks:
        print_warning("No hosts found. Cannot remove anything.")
        return

    # We'll display them in a table
    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")

        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"))

    # Prompt which host to remove
    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_block = None
    if choice.isdigit():
        idx = int(choice)
        if idx < 1 or idx > len(all_blocks):
            print_warning(f"Row number {idx} is invalid.")
            return
        target_block = all_blocks[idx - 1]
    else:
        # They 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

    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:
        print_warning("Target block has no Host label. Cannot remove.")
        return

    # Check IdentityFile for old pub key
    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:
                    # This is the EXACT line that might appear in authorized_keys
                    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)

    # Finally, remove the local subdirectory
    host_dir = os.path.join(conf_dir, host_label)
    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:
                import shutil
                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 host folder {host_dir} not found. Nothing to remove.")

async def remove_old_key_remote(old_pub_data, user, hostname, port):
    """
    Remove the old key from remote authorized_keys if found, matching EXACT lines.
    We'll do:
      grep -Fxq  for existence
      grep -vFx  to remove it
    """
    # 1) Check if old_pub_data is present in authorized_keys (exact line)
    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

    # 2) Actually remove it by ignoring EXACT matches to that line
    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).")