refactor
This commit is contained in:
parent
414266eefc
commit
ffb1f7e204
5 changed files with 636 additions and 312 deletions
252
add_host.py
252
add_host.py
|
@ -2,109 +2,211 @@
|
|||
|
||||
import os
|
||||
import subprocess
|
||||
import glob
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Set, Tuple
|
||||
from dataclasses import dataclass
|
||||
from .utils import print_error, print_warning, print_info, safe_input
|
||||
from .config import CONF_DIR
|
||||
|
||||
def add_host(conf_dir):
|
||||
@dataclass
|
||||
class HostConfig:
|
||||
label: str
|
||||
hostname: str
|
||||
user: str = "root"
|
||||
port: str = "22"
|
||||
identity_file: str = ""
|
||||
|
||||
def to_config_lines(self) -> List[str]:
|
||||
"""Convert host configuration to SSH config file lines."""
|
||||
lines = [
|
||||
f"Host {self.label}",
|
||||
f" HostName {self.hostname}",
|
||||
f" User {self.user}",
|
||||
f" Port {self.port}"
|
||||
]
|
||||
if self.identity_file:
|
||||
lines.append(f" IdentityFile {self.identity_file}")
|
||||
return lines
|
||||
|
||||
def get_existing_hosts(conf_dir: str) -> Tuple[Set[str], Dict[str, str]]:
|
||||
"""
|
||||
Get existing host labels and hostnames from config files.
|
||||
Returns (host_labels, hostname_to_label) where:
|
||||
- host_labels is a set of existing host labels
|
||||
- hostname_to_label maps hostnames to their labels
|
||||
"""
|
||||
host_labels = set()
|
||||
hostname_to_label = {}
|
||||
|
||||
pattern = os.path.join(conf_dir, "*", "config")
|
||||
for config_file in glob.glob(pattern):
|
||||
try:
|
||||
with open(config_file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
current_label = None
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line or line.startswith('#'):
|
||||
continue
|
||||
|
||||
if line.lower().startswith('host '):
|
||||
labels = line.split()[1:]
|
||||
for label in labels:
|
||||
if '*' not in label:
|
||||
current_label = label
|
||||
host_labels.add(label)
|
||||
break
|
||||
elif current_label and line.lower().startswith('hostname '):
|
||||
hostname = line.split(None, 1)[1].strip()
|
||||
hostname_to_label[hostname] = current_label
|
||||
|
||||
except Exception as e:
|
||||
print_warning(f"Error reading config file {config_file}: {e}")
|
||||
continue
|
||||
|
||||
return host_labels, hostname_to_label
|
||||
|
||||
def generate_ssh_key(key_path: Path) -> bool:
|
||||
"""Generate a new ED25519 SSH key pair."""
|
||||
try:
|
||||
subprocess.check_call([
|
||||
"ssh-keygen",
|
||||
"-q",
|
||||
"-t", "ed25519",
|
||||
"-N", "",
|
||||
"-f", str(key_path)
|
||||
])
|
||||
print_info(f"Generated new SSH key at {key_path}")
|
||||
return True
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error generating SSH key: {e}")
|
||||
return False
|
||||
|
||||
def copy_ssh_key(key_path: Path, user: str, hostname: str, port: str) -> bool:
|
||||
"""Copy SSH public key to remote server."""
|
||||
try:
|
||||
cmd = ["ssh-copy-id", "-i", str(key_path)]
|
||||
if port != "22":
|
||||
cmd.extend(["-p", port])
|
||||
cmd.append(f"{user}@{hostname}")
|
||||
|
||||
subprocess.check_call(cmd)
|
||||
print_info("Key successfully copied to remote server.")
|
||||
return True
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error copying key to server: {e}")
|
||||
return False
|
||||
|
||||
def write_config_file(config_path: Path, config_lines: List[str]) -> bool:
|
||||
"""Write SSH config lines to file."""
|
||||
try:
|
||||
config_path.write_text("\n".join(config_lines) + "\n")
|
||||
print_info(f"Created/updated config at: {config_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print_error(f"Failed to write config to {config_path}: {e}")
|
||||
return False
|
||||
|
||||
def add_host(conf_dir: str) -> bool:
|
||||
"""
|
||||
Interactive prompt to create a new SSH host in ~/.ssh/conf/<label>/config.
|
||||
Offers to generate a new SSH key pair (ed25519) quietly (-q),
|
||||
and then prompt to copy that key to the remote server via ssh-copy-id.
|
||||
|
||||
Returns True if host was added successfully, False otherwise.
|
||||
"""
|
||||
print_info("Adding a new SSH host...")
|
||||
|
||||
host_label = safe_input("Enter Host label (e.g. myserver): ")
|
||||
if host_label is None:
|
||||
return # User canceled (Ctrl+C)
|
||||
host_label = host_label.strip()
|
||||
if not host_label:
|
||||
print_error("Host label cannot be empty.")
|
||||
return
|
||||
# Get existing hosts to check for duplicates
|
||||
existing_labels, hostname_to_label = get_existing_hosts(conf_dir)
|
||||
|
||||
hostname = safe_input("Enter HostName (IP or domain): ")
|
||||
if hostname is None:
|
||||
return
|
||||
hostname = hostname.strip()
|
||||
if not hostname:
|
||||
print_error("HostName cannot be empty.")
|
||||
return
|
||||
# Get host label
|
||||
while True:
|
||||
host_label = safe_input("Enter Host label (e.g. myserver): ")
|
||||
if not host_label or host_label is None:
|
||||
print_error("Host label cannot be empty.")
|
||||
return False
|
||||
|
||||
host_label = host_label.strip()
|
||||
if host_label in existing_labels:
|
||||
print_error(f"A host with label '{host_label}' already exists. Please choose a different label.")
|
||||
continue
|
||||
break
|
||||
|
||||
# Get hostname
|
||||
while True:
|
||||
hostname = safe_input("Enter HostName (IP or domain): ")
|
||||
if not hostname or hostname is None:
|
||||
print_error("HostName cannot be empty.")
|
||||
return False
|
||||
|
||||
hostname = hostname.strip()
|
||||
if hostname in hostname_to_label:
|
||||
existing_label = hostname_to_label[hostname]
|
||||
print_error(f"This hostname is already configured for host '{existing_label}'. Please use a different hostname.")
|
||||
continue
|
||||
break
|
||||
|
||||
user = safe_input("Enter username (default: 'root'): ")
|
||||
# Get optional parameters
|
||||
user = safe_input("Enter username (default: 'root'): ") or "root"
|
||||
if user is None:
|
||||
return
|
||||
user = user.strip() or "root"
|
||||
|
||||
port = safe_input("Enter SSH port (default: 22): ")
|
||||
return False
|
||||
|
||||
port = safe_input("Enter SSH port (default: 22): ") or "22"
|
||||
if port is None:
|
||||
return
|
||||
port = port.strip() or "22"
|
||||
return False
|
||||
|
||||
# Create subdirectory: ~/.ssh/conf/<label>
|
||||
host_dir = os.path.join(conf_dir, host_label)
|
||||
if os.path.exists(host_dir):
|
||||
# Create host configuration
|
||||
host_config = HostConfig(
|
||||
label=host_label,
|
||||
hostname=hostname,
|
||||
user=user.strip(),
|
||||
port=port.strip()
|
||||
)
|
||||
|
||||
# Setup directory structure
|
||||
host_dir = Path(conf_dir) / host_label
|
||||
if host_dir.exists():
|
||||
print_warning(f"Directory {host_dir} already exists; continuing anyway.")
|
||||
else:
|
||||
os.makedirs(host_dir, mode=0o700, exist_ok=True)
|
||||
print_info(f"Created directory: {host_dir}")
|
||||
host_dir.mkdir(mode=0o700, exist_ok=True)
|
||||
print_info(f"Created directory: {host_dir}")
|
||||
|
||||
config_path = os.path.join(host_dir, "config")
|
||||
if os.path.exists(config_path):
|
||||
config_path = host_dir / "config"
|
||||
if config_path.exists():
|
||||
print_warning(f"Config file already exists: {config_path}; it will be overwritten/updated.")
|
||||
|
||||
# Handle SSH key generation
|
||||
gen_key_choice = safe_input("Generate a new ed25519 SSH key for this host? (y/n): ")
|
||||
if gen_key_choice is None:
|
||||
return
|
||||
gen_key_choice = gen_key_choice.lower().strip()
|
||||
return False
|
||||
|
||||
identity_file = ""
|
||||
if gen_key_choice == 'y':
|
||||
key_path = os.path.join(host_dir, "id_ed25519")
|
||||
if os.path.exists(key_path):
|
||||
if gen_key_choice.lower().strip() == 'y':
|
||||
key_path = host_dir / "id_ed25519"
|
||||
|
||||
if key_path.exists():
|
||||
print_warning(f"{key_path} already exists. Skipping generation.")
|
||||
identity_file = key_path
|
||||
host_config.identity_file = str(key_path)
|
||||
else:
|
||||
cmd = ["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-f", key_path]
|
||||
try:
|
||||
subprocess.check_call(cmd)
|
||||
print_info(f"Generated new SSH key at {key_path}")
|
||||
identity_file = key_path
|
||||
|
||||
if generate_ssh_key(key_path):
|
||||
host_config.identity_file = str(key_path)
|
||||
|
||||
# Prompt to copy the key
|
||||
copy_key = safe_input("Would you like to copy this key to the server now? (y/n): ")
|
||||
if copy_key is None:
|
||||
return
|
||||
return False
|
||||
|
||||
if copy_key.lower().strip() == 'y':
|
||||
ssh_copy_cmd = ["ssh-copy-id", "-i", key_path]
|
||||
if port != "22":
|
||||
ssh_copy_cmd += ["-p", port]
|
||||
ssh_copy_cmd.append(f"{user}@{hostname}")
|
||||
try:
|
||||
subprocess.check_call(ssh_copy_cmd)
|
||||
print_info("Key successfully copied to remote server.")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error copying key to server: {e}")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print_error(f"Error generating SSH key: {e}")
|
||||
copy_ssh_key(key_path, host_config.user, host_config.hostname, host_config.port)
|
||||
else:
|
||||
# Handle existing key
|
||||
existing_key = safe_input("Enter existing IdentityFile path (or leave empty to skip): ")
|
||||
if existing_key is None:
|
||||
return
|
||||
existing_key = existing_key.strip()
|
||||
if existing_key:
|
||||
identity_file = os.path.expanduser(existing_key)
|
||||
return False
|
||||
|
||||
if existing_key.strip():
|
||||
host_config.identity_file = os.path.expanduser(existing_key.strip())
|
||||
|
||||
config_lines = [
|
||||
f"Host {host_label}",
|
||||
f" HostName {hostname}",
|
||||
f" User {user}",
|
||||
f" Port {port}"
|
||||
]
|
||||
if identity_file:
|
||||
config_lines.append(f" IdentityFile {identity_file}")
|
||||
|
||||
try:
|
||||
with open(config_path, "w") as f:
|
||||
for line in config_lines:
|
||||
f.write(line + "\n")
|
||||
print_info(f"Created/updated config at: {config_path}")
|
||||
except Exception as e:
|
||||
print_error(f"Failed to write config to {config_path}: {e}")
|
||||
# Write the config file
|
||||
return write_config_file(config_path, host_config.to_config_lines())
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue