refactor
This commit is contained in:
parent
414266eefc
commit
ffb1f7e204
5 changed files with 636 additions and 312 deletions
388
regen_key.py
388
regen_key.py
|
|
@ -1,6 +1,8 @@
|
|||
import os
|
||||
import subprocess
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Tuple, Any
|
||||
from .utils import (
|
||||
print_info,
|
||||
print_warning,
|
||||
|
|
@ -14,105 +16,106 @@ from .list_hosts import (
|
|||
sort_by_ip
|
||||
)
|
||||
|
||||
async def regenerate_key(conf_dir):
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
def validate_key_path(key_path: str) -> bool:
|
||||
"""Validate that the key path and its directory are valid."""
|
||||
try:
|
||||
key_path = Path(key_path)
|
||||
key_dir = key_path.parent
|
||||
|
||||
if not key_dir.exists():
|
||||
key_dir.mkdir(mode=0o700, parents=True)
|
||||
elif not os.access(str(key_dir), os.W_OK):
|
||||
print_error(f"No write permission for directory: {key_dir}")
|
||||
return False
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print_error(f"Error validating key path: {e}")
|
||||
return False
|
||||
|
||||
print_info("Regenerate Key - Step 1: Show current hosts...\n")
|
||||
def generate_new_key(key_path: str, user: str, hostname: str, port: int) -> bool:
|
||||
"""Generate a new ED25519 SSH key and optionally copy it to the remote server."""
|
||||
if not validate_key_path(key_path):
|
||||
return False
|
||||
|
||||
# 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
|
||||
print_info("Generating new ed25519 SSH key...")
|
||||
try:
|
||||
subprocess.check_call([
|
||||
"ssh-keygen",
|
||||
"-q",
|
||||
"-t", "ed25519",
|
||||
"-N", "",
|
||||
"-f", key_path
|
||||
])
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error generating new SSH key: {e}")
|
||||
return False
|
||||
|
||||
print_info(f"Generated new SSH key at {key_path}")
|
||||
|
||||
# Ask to copy the key
|
||||
copy_choice = safe_input("Copy new key to remote now? (y/n): ")
|
||||
if not copy_choice or not copy_choice.lower().startswith('y'):
|
||||
return True
|
||||
|
||||
from tabulate import tabulate
|
||||
print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
|
||||
print(tabulate(final_data, headers=headers, tablefmt="grid"))
|
||||
try:
|
||||
ssh_copy_cmd = ["ssh-copy-id", "-i", key_path]
|
||||
if port != 22:
|
||||
ssh_copy_cmd += ["-p", str(port)]
|
||||
ssh_copy_cmd.append(f"{user}@{hostname}")
|
||||
|
||||
subprocess.check_call(ssh_copy_cmd)
|
||||
print_info("New key successfully copied to remote server.")
|
||||
return True
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error copying new key: {e}")
|
||||
return False
|
||||
|
||||
# 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))
|
||||
def update_config_with_key(config_path: Path, new_key_path: str) -> bool:
|
||||
"""Update the SSH config file with the new identity file."""
|
||||
try:
|
||||
with open(config_path, 'r') as f:
|
||||
config_lines = [
|
||||
line.rstrip('\n')
|
||||
for line in f
|
||||
if not line.strip().lower().startswith('identityfile')
|
||||
]
|
||||
|
||||
config_lines.append(f" IdentityFile {new_key_path}")
|
||||
|
||||
with open(config_path, 'w') as f:
|
||||
f.write('\n'.join(config_lines) + '\n')
|
||||
|
||||
print_info(f"Updated config file with new IdentityFile: {new_key_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print_error(f"Failed to update config file: {e}")
|
||||
return 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)
|
||||
|
||||
choice = safe_input("Enter the row number or the Host label to regenerate: ")
|
||||
if choice is None:
|
||||
return
|
||||
choice = choice.strip()
|
||||
def find_target_host(sorted_rows: List[Tuple], choice: str) -> Optional[Tuple]:
|
||||
"""Find the target host based on user choice."""
|
||||
if not choice:
|
||||
print_error("No choice given.")
|
||||
return
|
||||
|
||||
target_tuple = None
|
||||
return None
|
||||
|
||||
if choice.isdigit():
|
||||
row_idx = int(choice)
|
||||
if row_idx < 1 or row_idx > len(sorted_rows):
|
||||
print_warning(f"Invalid row number {row_idx}.")
|
||||
return
|
||||
target_tuple = sorted_rows[row_idx - 1]
|
||||
else:
|
||||
# user typed a label
|
||||
for t in sorted_rows:
|
||||
if t[0] == choice: # t[0] is host_label
|
||||
target_tuple = t
|
||||
break
|
||||
if not target_tuple:
|
||||
print_warning(f"No matching host label '{choice}' found.")
|
||||
return
|
||||
return None
|
||||
return sorted_rows[row_idx - 1]
|
||||
|
||||
# User typed a label
|
||||
for row in sorted_rows:
|
||||
if row[0] == choice: # row[0] is host_label
|
||||
return row
|
||||
|
||||
print_warning(f"No matching host label '{choice}' found.")
|
||||
return None
|
||||
|
||||
# 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]
|
||||
|
||||
# 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
|
||||
|
||||
# 4) Remove old local key files
|
||||
expanded_key = os.path.expanduser(identity_file)
|
||||
pub_path = expanded_key + ".pub"
|
||||
old_pub_data = ""
|
||||
|
||||
if os.path.isfile(pub_path):
|
||||
try:
|
||||
with open(pub_path, "r") as f:
|
||||
old_pub_data = f.read().rstrip("\n")
|
||||
except Exception as e:
|
||||
print_warning(f"Could not read old pub key: {e}")
|
||||
|
||||
print_info("Removing old key files locally...")
|
||||
for path in [expanded_key, pub_path]:
|
||||
def remove_key_files(key_paths: List[str]) -> None:
|
||||
"""Remove SSH key files."""
|
||||
for path in key_paths:
|
||||
if os.path.isfile(path):
|
||||
try:
|
||||
os.remove(path)
|
||||
|
|
@ -120,67 +123,196 @@ async def regenerate_key(conf_dir):
|
|||
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
|
||||
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
|
||||
async def regenerate_key(conf_dir: str) -> bool:
|
||||
"""
|
||||
Regenerate the SSH key for a chosen host by:
|
||||
1) Displaying the unified table of hosts
|
||||
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
|
||||
|
||||
Returns True if key was successfully regenerated, False otherwise.
|
||||
"""
|
||||
print_info("Regenerate Key - Step 1: Show current hosts...\n")
|
||||
|
||||
# 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]
|
||||
if port != 22:
|
||||
ssh_copy_cmd += ["-p", str(port)]
|
||||
ssh_copy_cmd.append(f"{user}@{hostname}")
|
||||
# Get host list
|
||||
headers, final_data = await build_host_list_table(conf_dir)
|
||||
if not final_data:
|
||||
print_warning("No hosts found. Cannot regenerate a key.")
|
||||
return False
|
||||
|
||||
# Display host table
|
||||
from tabulate import tabulate
|
||||
print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
|
||||
print(tabulate(final_data, headers=headers, tablefmt="grid"))
|
||||
|
||||
# Get host blocks and sort them
|
||||
all_blocks = []
|
||||
pattern = os.path.join(conf_dir, "*", "config")
|
||||
for cfile in glob.glob(pattern):
|
||||
all_blocks.extend(load_config_file(cfile))
|
||||
|
||||
results = await gather_host_info(all_blocks)
|
||||
sorted_rows = sort_by_ip(results)
|
||||
|
||||
# Get user choice
|
||||
choice = safe_input("Enter the row number or the Host label to regenerate: ")
|
||||
if choice is None:
|
||||
return False
|
||||
|
||||
target_tuple = find_target_host(sorted_rows, choice.strip())
|
||||
if not target_tuple:
|
||||
return False
|
||||
|
||||
# Get host information
|
||||
host_label, user, _, hostname, *_ = target_tuple
|
||||
|
||||
# Find config block
|
||||
found_block = next(
|
||||
(b for b in all_blocks if b.get("Host") == host_label),
|
||||
None
|
||||
)
|
||||
if not found_block:
|
||||
print_warning(f"No config block found for '{host_label}'.")
|
||||
return False
|
||||
|
||||
port = int(found_block.get("Port", "22"))
|
||||
identity_file = found_block.get("IdentityFile", "")
|
||||
|
||||
# Handle missing identity file
|
||||
if not identity_file or identity_file == "N/A":
|
||||
print_warning("No existing SSH key found in the configuration.")
|
||||
gen_choice = safe_input("Would you like to generate a new key? (y/n): ")
|
||||
if not gen_choice or not gen_choice.lower().startswith('y'):
|
||||
return False
|
||||
|
||||
# Set up new key path and generate key
|
||||
host_dir = Path(conf_dir) / host_label
|
||||
host_dir.mkdir(mode=0o700, exist_ok=True)
|
||||
new_key_path = str(host_dir / "id_ed25519")
|
||||
|
||||
if not generate_new_key(new_key_path, user, hostname, port):
|
||||
return False
|
||||
|
||||
# Update config with new key
|
||||
config_path = host_dir / "config"
|
||||
return update_config_with_key(config_path, new_key_path)
|
||||
|
||||
# Handle existing key regeneration
|
||||
expanded_key = os.path.expanduser(identity_file)
|
||||
pub_path = expanded_key + ".pub"
|
||||
old_pub_data = ""
|
||||
|
||||
# Try to read old public key
|
||||
if os.path.isfile(pub_path):
|
||||
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}")
|
||||
old_pub_data = Path(pub_path).read_text().rstrip("\n")
|
||||
except Exception as e:
|
||||
print_warning(f"Could not read old pub key: {e}")
|
||||
|
||||
# 7) Remove old key from authorized_keys if we had old_pub_data
|
||||
# Remove old key files
|
||||
print_info("Removing old key files locally...")
|
||||
remove_key_files([expanded_key, pub_path])
|
||||
|
||||
# Generate new key
|
||||
if not generate_new_key(expanded_key, user, hostname, port):
|
||||
return False
|
||||
|
||||
# Remove old key from remote if we had it
|
||||
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.")
|
||||
|
||||
return True
|
||||
|
||||
async def remove_old_key_remote(old_pub_data, user, hostname, port):
|
||||
async def remove_old_key_remote(old_pub_data: str, user: str, hostname: str, port: int) -> bool:
|
||||
"""
|
||||
Checks and removes the exact matching line from remote authorized_keys
|
||||
via grep -Fxq / grep -vFx.
|
||||
Remove the old public key from remote authorized_keys file.
|
||||
Returns True if successful, False otherwise.
|
||||
"""
|
||||
check_cmd = [
|
||||
"ssh", "-o", "StrictHostKeyChecking=no",
|
||||
# Escape the public key data for shell safety
|
||||
escaped_key = old_pub_data.replace('"', '\\"')
|
||||
|
||||
# First check if authorized_keys exists
|
||||
check_file_cmd = [
|
||||
"ssh",
|
||||
"-o", "StrictHostKeyChecking=no",
|
||||
"-p", str(port),
|
||||
f"{user}@{hostname}",
|
||||
f'grep -Fxq "{old_pub_data}" ~/.ssh/authorized_keys'
|
||||
"test -f ~/.ssh/authorized_keys && echo 'exists'"
|
||||
]
|
||||
found_key = False
|
||||
|
||||
try:
|
||||
subprocess.check_call(check_cmd)
|
||||
found_key = True
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
result = subprocess.run(check_file_cmd, capture_output=True, text=True)
|
||||
if result.returncode != 0 or 'exists' not in result.stdout:
|
||||
print_warning("No authorized_keys file found on remote host.")
|
||||
return False
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error checking authorized_keys: {e}")
|
||||
return False
|
||||
|
||||
if not found_key:
|
||||
print_warning("No old key found on remote authorized_keys.")
|
||||
return
|
||||
# Create a temporary file for the sed script
|
||||
temp_script = """#!/bin/bash
|
||||
set -e
|
||||
KEYS_FILE="$HOME/.ssh/authorized_keys"
|
||||
TEMP_FILE="$HOME/.ssh/authorized_keys.tmp"
|
||||
grep -Fxv "$1" "$KEYS_FILE" > "$TEMP_FILE"
|
||||
mv "$TEMP_FILE" "$KEYS_FILE"
|
||||
chmod 600 "$KEYS_FILE"
|
||||
"""
|
||||
|
||||
# Create a temporary script on the remote host
|
||||
setup_cmd = [
|
||||
"ssh",
|
||||
"-o", "StrictHostKeyChecking=no",
|
||||
"-p", str(port),
|
||||
f"{user}@{hostname}",
|
||||
f'cat > ~/.ssh/remove_key.sh << \'EOF\'\n{temp_script}\nEOF\n'
|
||||
f'chmod +x ~/.ssh/remove_key.sh'
|
||||
]
|
||||
|
||||
try:
|
||||
subprocess.run(setup_cmd, check=True, capture_output=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Failed to create temporary script: {e}")
|
||||
return False
|
||||
|
||||
# Execute the script with the key
|
||||
remove_cmd = [
|
||||
"ssh", "-o", "StrictHostKeyChecking=no",
|
||||
"ssh",
|
||||
"-o", "StrictHostKeyChecking=no",
|
||||
"-p", str(port),
|
||||
f"{user}@{hostname}",
|
||||
f'grep -vFx "{old_pub_data}" ~/.ssh/authorized_keys > ~/.ssh/tmp && mv ~/.ssh/tmp ~/.ssh/authorized_keys'
|
||||
f'~/.ssh/remove_key.sh "{escaped_key}"'
|
||||
]
|
||||
|
||||
try:
|
||||
subprocess.check_call(remove_cmd)
|
||||
subprocess.run(remove_cmd, check=True, capture_output=True)
|
||||
print_info("Old public key removed from remote authorized_keys.")
|
||||
except subprocess.CalledProcessError:
|
||||
print_warning("Failed to remove old key from remote authorized_keys (permissions or other error).")
|
||||
|
||||
# Clean up the temporary script
|
||||
cleanup_cmd = [
|
||||
"ssh",
|
||||
"-o", "StrictHostKeyChecking=no",
|
||||
"-p", str(port),
|
||||
f"{user}@{hostname}",
|
||||
"rm -f ~/.ssh/remove_key.sh"
|
||||
]
|
||||
subprocess.run(cleanup_cmd, check=True, capture_output=True)
|
||||
|
||||
return True
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Failed to remove old key: {e}")
|
||||
# Try to clean up even if removal failed
|
||||
try:
|
||||
subprocess.run(cleanup_cmd, check=True, capture_output=True)
|
||||
except:
|
||||
pass
|
||||
return False
|
||||
except Exception as e:
|
||||
print_error(f"Unexpected error removing old key: {e}")
|
||||
return False
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue