added additional functions

This commit is contained in:
Arctic 2025-03-07 06:08:30 -06:00
parent 1e55abf811
commit 67a5534246
8 changed files with 654 additions and 107 deletions

View file

@ -2,7 +2,7 @@
import os
import subprocess
from .utils import print_error, print_warning, print_info
from .utils import print_error, print_warning, print_info, safe_input
def add_host(conf_dir):
"""
@ -12,18 +12,31 @@ def add_host(conf_dir):
"""
print_info("Adding a new SSH host...")
host_label = input("Enter Host label (e.g. myserver): ").strip()
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
hostname = input("Enter HostName (IP or domain): ").strip()
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
user = input("Enter username (default: 'root'): ").strip() or "root"
port = input("Enter SSH port (default: 22): ").strip() or "22"
user = safe_input("Enter username (default: 'root'): ")
if user is None:
return
user = user.strip() or "root"
port = safe_input("Enter SSH port (default: 22): ")
if port is None:
return
port = port.strip() or "22"
# Create subdirectory: ~/.ssh/conf/<label>
host_dir = os.path.join(conf_dir, host_label)
@ -37,26 +50,29 @@ def add_host(conf_dir):
if os.path.exists(config_path):
print_warning(f"Config file already exists: {config_path}; it will be overwritten/updated.")
# Ask about generating an SSH key
gen_key_choice = input("Generate a new ed25519 SSH key for this host? (y/n): ").lower().strip()
identity_file = ""
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()
identity_file = ""
if gen_key_choice == 'y':
key_path = os.path.join(host_dir, "id_ed25519")
if os.path.exists(key_path):
print_warning(f"{key_path} already exists. Skipping generation.")
identity_file = key_path
else:
# Generate a new SSH key (quietly, suppressing randomart)
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
# Prompt to copy the key to the server
copy_key = input("Would you like to copy this key to the server now? (y/n): ").lower().strip()
if copy_key == 'y':
# 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
if copy_key.lower().strip() == 'y':
ssh_copy_cmd = ["ssh-copy-id", "-i", key_path]
if port != "22":
ssh_copy_cmd += ["-p", port]
@ -69,12 +85,13 @@ def add_host(conf_dir):
except subprocess.CalledProcessError as e:
print_error(f"Error generating SSH key: {e}")
else:
# If not generating a new key, optionally ask for an existing path
existing_key = input("Enter existing IdentityFile path (or leave empty to skip): ").strip()
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)
# Build the config lines
config_lines = [
f"Host {host_label}",
f" HostName {hostname}",
@ -84,7 +101,6 @@ def add_host(conf_dir):
if identity_file:
config_lines.append(f" IdentityFile {identity_file}")
# Write (or overwrite) the config
try:
with open(config_path, "w") as f:
for line in config_lines: