SSH-key-Manager/list_hosts.py

# ssh_manager/list_hosts.py

import os
import glob
import socket
import asyncio
import ipaddress
from tabulate import tabulate
from collections import OrderedDict

from .utils import print_warning, print_error, Colors

async def check_ssh_port(ip_address, port):
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(ip_address, port), timeout=1
        )
        writer.close()
        await writer.wait_closed()
        return True
    except:
        return False

def load_config_file(file_path):
    blocks = []
    host_data = None
    try:
        with open(file_path, 'r') as f:
            lines = f.readlines()
    except Exception as e:
        print_error(f"Error reading SSH config file {file_path}: {e}")
        return blocks

    for line in lines:
        stripped_line = line.strip()
        if not stripped_line or stripped_line.startswith('#'):
            continue

        if stripped_line.lower().startswith('host '):
            host_labels = stripped_line.split()[1:]
            for label in host_labels:
                if '*' not in label:
                    if host_data:
                        blocks.append(host_data)
                    host_data = OrderedDict({'Host': label})
                    break
        elif host_data:
            if ' ' in stripped_line:
                key, value = stripped_line.split(None, 1)
                host_data[key] = value.strip()

    if host_data:
        blocks.append(host_data)
    return blocks

async def gather_host_info(all_host_blocks):
    """
    Given a list of host blocks, gather full info:
      - resolved IP
      - port open check
      - 'Conf Directory' coloring if IdentityFile != 'N/A'
    Returns a list of 7-tuples:
      (host_label, user, colored_port, hostname, colored_ip, conf_path_display, raw_ip)
    """
    results = []

    async def process_block(h):
        host_label    = h.get('Host', 'N/A')
        hostname      = h.get('HostName', 'N/A')
        user          = h.get('User', 'N/A')
        port          = int(h.get('Port', '22'))
        identity_file = h.get('IdentityFile', 'N/A')

        # Resolve IP
        try:
            raw_ip = socket.gethostbyname(hostname)
            colored_ip = f"{Colors.GREEN}{raw_ip}{Colors.RESET}"
        except socket.error:
            raw_ip = "N/A"
            colored_ip = f"{Colors.RED}N/A{Colors.RESET}"

        # Check port
        if raw_ip != "N/A":
            port_open = await check_ssh_port(raw_ip, port)
            colored_port = (
                f"{Colors.GREEN}{port}{Colors.RESET}" if port_open else f"{Colors.RED}{port}{Colors.RESET}"
            )
        else:
            colored_port = f"{Colors.RED}{port}{Colors.RESET}"

        # Conf Directory = ~/.ssh/conf/<host_label>
        conf_path = f"~/.ssh/conf/{host_label}"
        # If there's an IdentityFile, color the conf path
        if identity_file != 'N/A':
            conf_path_display = f"{Colors.GREEN}{conf_path}{Colors.RESET}"
        else:
            conf_path_display = conf_path

        return (
            host_label,
            user,
            colored_port,
            hostname,
            colored_ip,
            conf_path_display,
            raw_ip  # uncolored IP for sorting
        )

    # Process blocks concurrently
    tasks = [process_block(b) for b in all_host_blocks]
    results = await asyncio.gather(*tasks)
    return results

def parse_ip(ip_str):
    """
    Convert a string IP to an ipaddress object for sorting.
    Returns None if invalid or 'N/A'.
    """
    import ipaddress
    try:
        return ipaddress.ip_address(ip_str)
    except ValueError:
        return None

def sort_by_ip(results):
    """
    Sort the 7-tuples by IP ascending, with 'N/A' last.
    """
    sortable = []
    for row in results:
        raw_ip = row[-1]
        ip_obj = parse_ip(raw_ip)
        sortable.append(((ip_obj is None, ip_obj), row))

    sortable.sort(key=lambda x: x[0])
    return [row for (_, row) in sortable]

async def build_host_list_table(conf_dir):
    """
    Gathers + sorts all hosts in conf_dir by IP ascending.
    Returns (headers, final_table_rows), each row omitting the raw_ip.
    """
    pattern    = os.path.join(conf_dir, "*", "config")
    conf_files = sorted(glob.glob(pattern))

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

    headers = ["No.", "Host", "User", "Port", "HostName", "IP Address", "Conf Directory"]
    if not all_host_blocks:
        return headers, []

    results = await gather_host_info(all_host_blocks)
    sorted_rows = sort_by_ip(results)

    # Build final table
    final_data = []
    for idx, row in enumerate(sorted_rows, start=1):
        # row is (host_label, user, colored_port, hostname, colored_ip, conf_path_display, raw_ip)
        final_data.append([idx] + list(row[:-1]))

    return headers, final_data

async def list_hosts(conf_dir):
    headers, final_data = await build_host_list_table(conf_dir)
    if not final_data:
        print_warning("No hosts found. The server list is empty.")
        print("\nSSH Conf Subdirectory Host List")
        from tabulate import tabulate
        print(tabulate([], headers=headers, tablefmt="grid"))
        return

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