refactored code

regen_key.py

@@ -1,88 +1,54 @@
 import os
-import glob
 import subprocess
 import asyncio
-from collections import OrderedDict
 from .utils import (
     print_info,
     print_warning,
     print_error,
-    safe_input,
-    Colors
+    safe_input
+)
+from .list_hosts import (
+    build_host_list_table,
+    load_config_file,
+    gather_host_info,
+    sort_by_ip
 )
-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)
-
-    if not all_blocks:
-        return []
-    return all_blocks
 
 async def regenerate_key(conf_dir):
     """
-    Menu-driven function to regenerate a key for a selected host:
-      1) Show the host table with row numbers
-      2) Let user pick row # or label
-      3) Read old pub data
-      4) Remove old local key files
-      5) Generate new key
-      6) Copy new key
-      7) Remove old key from remote authorized_keys (improved logic to detect existence)
+    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) Gather host blocks
-    all_blocks = await get_all_host_blocks(conf_dir)
-    if not all_blocks:
+    # 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
 
-    # Display them in a table (similar to list_hosts):
-    import socket
     from tabulate import tabulate
+    print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
+    print(tabulate(final_data, headers=headers, tablefmt="grid"))
 
-    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")
+    # 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))
 
-        # Check IP
-        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
+    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)
 
-        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"))
-
-    # 2) Prompt for row # or label
     choice = safe_input("Enter the row number or the Host label to regenerate: ")
     if choice is None:
         return
@@ -91,54 +57,60 @@ async def regenerate_key(conf_dir):
         print_error("No choice given.")
         return
 
-    target_block = None
+    target_tuple = None
     if choice.isdigit():
         row_idx = int(choice)
-        if row_idx < 1 or row_idx > len(all_blocks):
+        if row_idx < 1 or row_idx > len(sorted_rows):
             print_warning(f"Invalid row number {row_idx}.")
             return
-        target_block = all_blocks[row_idx - 1]
+        target_tuple = sorted_rows[row_idx - 1]
     else:
         # user typed a label
-        for b in all_blocks:
-            if b.get("Host") == choice:
-                target_block = b
+        for t in sorted_rows:
+            if t[0] == choice:  # t[0] is host_label
+                target_tuple = t
                 break
-        if not target_block:
+        if not target_tuple:
             print_warning(f"No matching host label '{choice}' found.")
             return
 
-    # 3) Gather info from block
-    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", "")
+    # 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]
 
-    if not host_label or not hostname:
-        print_error("Missing Host or HostName; cannot regenerate.")
+    # 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
 
-    # Derive local paths
+    # 4) Remove old local key files
     expanded_key  = os.path.expanduser(identity_file)
-    key_dir       = os.path.dirname(expanded_key)
     pub_path      = expanded_key + ".pub"
+    old_pub_data  = ""
 
-    # 3a) Read old pub key data
-    old_pub_data = ""
     if os.path.isfile(pub_path):
         try:
             with open(pub_path, "r") as f:
-                old_pub_data = f.read().strip()
+                old_pub_data = f.read().rstrip("\n")
         except Exception as e:
             print_warning(f"Could not read old pub key: {e}")
-    else:
-        print_warning("No old pub key found locally.")
 
-    # 4) Remove old local key files
     print_info("Removing old key files locally...")
     for path in [expanded_key, pub_path]:
         if os.path.isfile(path):
@@ -150,7 +122,7 @@ async def regenerate_key(conf_dir):
 
     # 5) Generate new key
     print_info("Generating new ed25519 SSH key...")
-    new_key_path = expanded_key  # Reuse the same path from config
+    new_key_path = expanded_key
     cmd = ["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-f", new_key_path]
     try:
         subprocess.check_call(cmd)
@@ -159,7 +131,7 @@ async def regenerate_key(conf_dir):
         print_error(f"Error generating new SSH key: {e}")
         return
 
-    # 6) Copy new key to remote
+    # 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]
@@ -172,7 +144,7 @@ async def regenerate_key(conf_dir):
         except subprocess.CalledProcessError as e:
             print_error(f"Error copying new key: {e}")
 
-    # 7) Remove old key from authorized_keys if old_pub_data is non-empty
+    # 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)
@@ -181,41 +153,34 @@ async def regenerate_key(conf_dir):
 
 async def remove_old_key_remote(old_pub_data, user, hostname, port):
     """
-    1) Check if old_pub_data is present on remote server (grep -q).
-    2) If found, remove it with grep -v ...
-    3) Otherwise, print "No old key found on remote."
+    Checks and removes the exact matching line from remote authorized_keys
+    via grep -Fxq / grep -vFx.
     """
-    # 1) Check if old_pub_data exists in authorized_keys
     check_cmd = [
-        "ssh",
-        "-o", "StrictHostKeyChecking=no",
+        "ssh", "-o", "StrictHostKeyChecking=no",
         "-p", str(port),
         f"{user}@{hostname}",
-        f"grep -F '{old_pub_data}' ~/.ssh/authorized_keys"
+        f'grep -Fxq "{old_pub_data}" ~/.ssh/authorized_keys'
     ]
     found_key = False
-
     try:
         subprocess.check_call(check_cmd)
         found_key = True
     except subprocess.CalledProcessError:
-        # grep returns exit code 1 if not found
         pass
 
     if not found_key:
         print_warning("No old key found on remote authorized_keys.")
         return
 
-    # 2) Actually remove it
     remove_cmd = [
-        "ssh",
-        "-o", "StrictHostKeyChecking=no",
+        "ssh", "-o", "StrictHostKeyChecking=no",
         "-p", str(port),
         f"{user}@{hostname}",
-        f"grep -v '{old_pub_data}' ~/.ssh/authorized_keys > ~/.ssh/tmp && mv ~/.ssh/tmp ~/.ssh/authorized_keys"
+        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 (permission or other error).")
+        print_warning("Failed to remove old key from remote authorized_keys (permissions or other error).")