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