added additional functions

regen_key.py

@@ -0,0 +1,221 @@
+import os
+import glob
+import subprocess
+import asyncio
+from collections import OrderedDict
+from .utils import (
+    print_info,
+    print_warning,
+    print_error,
+    safe_input,
+    Colors
+)
+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)
+    """
+    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:
+        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
+
+    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")
+
+        # 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
+
+        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
+    choice = choice.strip()
+    if not choice:
+        print_error("No choice given.")
+        return
+
+    target_block = None
+    if choice.isdigit():
+        row_idx = int(choice)
+        if row_idx < 1 or row_idx > len(all_blocks):
+            print_warning(f"Invalid row number {row_idx}.")
+            return
+        target_block = all_blocks[row_idx - 1]
+    else:
+        # user 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
+
+    # 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", "")
+
+    if not host_label or not hostname:
+        print_error("Missing Host or HostName; cannot regenerate.")
+        return
+    if not identity_file or identity_file == "N/A":
+        print_error("No IdentityFile found in config; cannot regenerate.")
+        return
+
+    # Derive local paths
+    expanded_key  = os.path.expanduser(identity_file)
+    key_dir       = os.path.dirname(expanded_key)
+    pub_path      = expanded_key + ".pub"
+
+    # 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()
+        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):
+            try:
+                os.remove(path)
+                print_info(f"Removed {path}")
+            except Exception as e:
+                print_warning(f"Could not remove {path}: {e}")
+
+    # 5) Generate new key
+    print_info("Generating new ed25519 SSH key...")
+    new_key_path = expanded_key  # Reuse the same path from config
+    cmd = ["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-f", new_key_path]
+    try:
+        subprocess.check_call(cmd)
+        print_info(f"Generated new SSH key at {new_key_path}")
+    except subprocess.CalledProcessError as e:
+        print_error(f"Error generating new SSH key: {e}")
+        return
+
+    # 6) Copy new key to remote
+    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]
+        if port != 22:
+            ssh_copy_cmd += ["-p", str(port)]
+        ssh_copy_cmd.append(f"{user}@{hostname}")
+        try:
+            subprocess.check_call(ssh_copy_cmd)
+            print_info("New key successfully copied to remote server.")
+        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
+    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)
+    else:
+        print_warning("No old pub key data found locally, skipping remote removal.")
+
+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."
+    """
+    # 1) Check if old_pub_data exists in authorized_keys
+    check_cmd = [
+        "ssh",
+        "-o", "StrictHostKeyChecking=no",
+        "-p", str(port),
+        f"{user}@{hostname}",
+        f"grep -F '{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",
+        "-p", str(port),
+        f"{user}@{hostname}",
+        f"grep -v '{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).")