refactored code

list_hosts.py

@@ -11,10 +11,6 @@ from collections import OrderedDict
 from .utils import print_warning, print_error, Colors
 
 async def check_ssh_port(ip_address, port):
-    """
-    Attempt to open an SSH connection to see if the port is open.
-    Returns True if successful, False otherwise.
-    """
     try:
         reader, writer = await asyncio.wait_for(
             asyncio.open_connection(ip_address, port), timeout=1
@@ -26,13 +22,8 @@ async def check_ssh_port(ip_address, port):
         return False
 
 def load_config_file(file_path):
-    """
-    Parse a single SSH config file and return a list of host blocks.
-    Each block is an OrderedDict with keys like 'Host', 'HostName', etc.
-    """
     blocks = []
     host_data = None
-
     try:
         with open(file_path, 'r') as f:
             lines = f.readlines()
@@ -42,11 +33,9 @@ def load_config_file(file_path):
 
     for line in lines:
         stripped_line = line.strip()
-        # Skip empty lines and comments
         if not stripped_line or stripped_line.startswith('#'):
             continue
 
-        # Start of a new Host block
         if stripped_line.lower().startswith('host '):
             host_labels = stripped_line.split()[1:]
             for label in host_labels:
@@ -64,64 +53,92 @@ def load_config_file(file_path):
         blocks.append(host_data)
     return blocks
 
-async def check_host(host):
+async def gather_host_info(all_host_blocks):
     """
-    Given a host block, resolve IP, check SSH port, etc.
-    Returns a tuple of:
-      1) Host label
-      2) User
-      3) Port (colored if open)
-      4) HostName
-      5) IP Address (colored if resolved)
-      6) Conf Directory (green if has IdentityFile, else no color)
-      7) raw_ip (uncolored string for sorting)
+    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)
     """
-    host_label    = host.get('Host', 'N/A')
-    hostname      = host.get('HostName', 'N/A')
-    user          = host.get('User', 'N/A')
-    port          = int(host.get('Port', '22'))
-    identity_file = host.get('IdentityFile', 'N/A')
+    results = []
 
-    # Resolve IP
-    try:
-        raw_ip = socket.gethostbyname(hostname)  # uncolored
-        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}"
+    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')
 
-    # 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}"
+        # 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
         )
-    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 green
-    if identity_file != 'N/A':
-        conf_path_display = f"{Colors.GREEN}{conf_path}{Colors.RESET}"
-    else:
-        conf_path_display = conf_path
+    # Process blocks concurrently
+    tasks = [process_block(b) for b in all_host_blocks]
+    results = await asyncio.gather(*tasks)
+    return results
 
-    # Return the data plus the uncolored IP for sorting
-    return (
-        host_label,
-        user,
-        colored_port,
-        hostname,
-        colored_ip,
-        conf_path_display,
-        raw_ip  # for sorting
-    )
-
-async def list_hosts(conf_dir):
+def parse_ip(ip_str):
     """
-    List out all hosts found in ~/.ssh/conf/*/config, sorted by IP in ascending order.
-    Columns: No., Host, User, Port, HostName, IP Address, Conf Directory
+    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))
@@ -133,42 +150,28 @@ async def list_hosts(conf_dir):
 
     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
 
-    # Gather full data for each host
-    tasks   = [check_host(h) for h in all_host_blocks]
-    results = await asyncio.gather(*tasks)
-
-    # We want to sort by IP ascending. results[i] is a tuple:
-    # (host_label, user, colored_port, hostname, colored_ip, conf_path, raw_ip)
-    # We'll parse raw_ip as an ipaddress for sorting. "N/A" => sort to the end.
-    def parse_ip(ip_str):
-        try:
-            return ipaddress.ip_address(ip_str)
-        except ValueError:
-            return None
-
-    # Convert the results into a list of (ip_obj, original_tuple)
-    # so we can sort, then rebuild the final data.
-    sortable = []
-    for row in results:
-        raw_ip = row[-1]  # last element
-        ip_obj = parse_ip(raw_ip)
-        # We'll sort None last by using a sort key that puts (True) after (False)
-        # e.g. (ip_obj is None, ip_obj)
-        sortable.append(((ip_obj is None, ip_obj), row))
-
-    # Sort by (is_none, ip_obj)
-    sortable.sort(key=lambda x: x[0])
-
-    # Rebuild the final display table, ignoring the raw_ip at the end
-    final_data = []
-    for idx, (_, row) in enumerate(sortable, start=1):
-        # row is (host_label, user, colored_port, hostname, colored_ip, conf_path, raw_ip)
-        final_data.append([idx] + list(row[:-1]))  # omit raw_ip
-
+    from tabulate import tabulate
     print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
     print(tabulate(final_data, headers=headers, tablefmt="grid"))