ssh-setup.sh

#!/usr/bin/env bash
# ============================================================================
# ssh-setup.sh — SSH Configuration Setup & Hardening Helper
# ============================================================================
#
# Description:
#   A defensive, idempotent bash script that simplifies SSH configuration
#   setup on Linux systems. Checks prerequisites, generates keys, manages
#   authorized_keys, and optionally hardens the SSH daemon configuration.
#
# Author:       SSH Setup Helper
# Version:      1.0.0
# License:      MIT
# Requirements: Bash 4.4+, Linux, root (for hardening mode only)
#
# Usage:
#   ./ssh-setup.sh [COMMAND] [OPTIONS]
#
# Commands:
#   all           Run full setup (check, generate key, configure)
#   check         Check SSH installation and configuration status
#   generate      Generate a new SSH key pair
#   add-key       Add a public key to authorized_keys
#   harden        Harden /etc/ssh/sshd_config (requires root)
#   status        Show current SSH configuration summary
#   help          Show this help message
#
# Options:
#   -e, --email <email>     Email label for generated key
#   -f, --file <path>       Path to key file (default: ~/.ssh/id_ed25519)
#   -k, --key <pubkey>      Public key string to add
#   -p, --port <port>       SSH port for hardening (default: 22)
#   -n, --no-backup         Skip backup before hardening (dangerous)
#   -y, --yes               Skip interactive prompts (non-interactive mode)
#   -d, --dry-run           Show what would be done without making changes
#   -v, --verbose           Enable verbose/debug output
#   -h, --help              Show this help message
#
# Examples:
#   ./ssh-setup.sh all
#   ./ssh-setup.sh generate --email admin@example.com
#   ./ssh-setup.sh add-key --key "ssh-ed25519 AAAA..."
#   ./ssh-setup.sh harden --port 2222 --yes
#   ./ssh-setup.sh status
#   ./ssh-setup.sh all --dry-run
#
# ============================================================================

# ---------------------------------------------------------------------------
# Strict Mode & Safety
# ---------------------------------------------------------------------------
set -Eeuo pipefail
shopt -s inherit_errexit 2>/dev/null || true
IFS=$'\n\t'

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
readonly SCRIPT_NAME="$(basename -- "${BASH_SOURCE[0]}")"
readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
readonly SCRIPT_VERSION="1.0.0"
readonly SSH_DIR="${HOME}/.ssh"
readonly DEFAULT_KEY_TYPE="ed25519"
readonly DEFAULT_KEY_PATH="${SSH_DIR}/id_ed25519"
readonly DEFAULT_SSHD_CONFIG="/etc/ssh/sshd_config"
readonly LOG_FILE="${SSH_DIR}/ssh-setup.log"
readonly DATE_FMT="%Y-%m-%d %H:%M:%S"
readonly BACKUP_SUFFIX=".bak.$(date +%Y%m%d%H%M%S)"

# Colors for terminal output
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly CYAN='\033[0;36m'
readonly BOLD='\033[1m'
readonly NC='\033[0m' # No Color

# ---------------------------------------------------------------------------
# Global Flags (set via CLI arguments)
# ---------------------------------------------------------------------------
COMMAND=""
VERBOSE=0
DRY_RUN=0
NON_INTERACTIVE=0
SKIP_BACKUP=0
USER_EMAIL=""
KEY_FILE=""
PUBLIC_KEY=""
SSH_PORT=22

# ---------------------------------------------------------------------------
# Logging Functions
# ---------------------------------------------------------------------------

# log LEVEL "message"
log() {
    local level="${1:?missing log level}"
    shift
    local msg="$*"
    local ts
    ts="$(date +"${DATE_FMT}")"
    local log_entry="[${ts}] [${level}] ${msg}"

    # Always write to log file if SSH_DIR exists
    if [[ -d "${SSH_DIR}" ]] || [[ "${DRY_RUN}" -eq 1 ]]; then
        printf '%s\n' "${log_entry}" >> "${LOG_FILE}" 2>/dev/null || true
    fi

    # Terminal output based on level
    case "${level}" in
        INFO)
            printf "${GREEN}[INFO]${NC} %s\n" "${msg}" >&2
            ;;
        WARN)
            printf "${YELLOW}[WARN]${NC} %s\n" "${msg}" >&2
            ;;
        ERROR)
            printf "${RED}[ERROR]${NC} %s\n" "${msg}" >&2
            ;;
        DEBUG)
            if [[ "${VERBOSE}" -eq 1 ]]; then
                printf "${CYAN}[DEBUG]${NC} %s\n" "${msg}" >&2
            fi
            ;;
        ACTION)
            printf "${BLUE}[ACTION]${NC} %s\n" "${msg}" >&2
            ;;
        *)
            printf '[%s] %s\n' "${level}" "${msg}" >&2
            ;;
    esac
}

log_info()    { log INFO "$@"; }
log_warn()    { log WARN "$@"; }
log_error()   { log ERROR "$@"; }
log_debug()   { log DEBUG "$@"; }
log_action()  { log ACTION "$@"; }

# ---------------------------------------------------------------------------
# Error Handling & Traps
# ---------------------------------------------------------------------------

cleanup() {
    local exit_code=$?
    if [[ ${exit_code} -ne 0 ]]; then
        log_error "Script exited with error code ${exit_code} at line ${BASH_LINENO[0]:-$LINENO}"
    fi
    exit "${exit_code}"
}

trap cleanup EXIT
trap 'log_error "Interrupted by signal"; exit 130' INT TERM
trap 'log_error "Error on line $LINENO, command: ${BASH_COMMAND}"' ERR

# ---------------------------------------------------------------------------
# Utility Functions
# ---------------------------------------------------------------------------

# require_cmd "command_name"
require_cmd() {
    local cmd="${1:?missing command name}"
    if ! command -v "${cmd}" &>/dev/null; then
        log_error "Required command not found: ${cmd}"
        return 1
    fi
    log_debug "Found required command: ${cmd}"
    return 0
}

# is_root — returns 0 if running as root
is_root() {
    [[ "${EUID:-$(id -u)}" -eq 0 ]]
}

# confirm "prompt message" — asks user for y/n, returns 0 for yes
confirm() {
    local prompt="${1:?missing prompt}"
    local answer

    if [[ "${NON_INTERACTIVE}" -eq 1 ]]; then
        log_info "Non-interactive mode: auto-confirming '${prompt}'"
        return 0
    fi

    while true; do
        printf "${BOLD}${prompt}${NC} [y/N]: " >&2
        read -r answer
        case "${answer,,}" in
            y|yes) return 0 ;;
            n|no|"") return 1 ;;
            *) printf "Please enter 'y' or 'n'.\n" >&2 ;;
        esac
    done
}

# dry_run_action "description" — logs action, skips if dry-run
dry_run_action() {
    local desc="${1:?missing description}"
    if [[ "${DRY_RUN}" -eq 1 ]]; then
        log_action "[DRY RUN] Would: ${desc}"
        return 0
    fi
    log_action "${desc}"
    return 1  # Caller should check return; 0 = skipped (dry run), 1 = execute
}

# ensure_dir PATH MODE — creates directory with specified permissions
ensure_dir() {
    local dir_path="${1:?missing directory path}"
    local mode="${2:?missing mode}"

    if [[ -d "${dir_path}" ]]; then
        log_debug "Directory already exists: ${dir_path}"
        # Verify permissions
        local current_mode
        current_mode="$(stat -c '%a' "${dir_path}" 2>/dev/null || stat -f '%Lp' "${dir_path}" 2>/dev/null)"
        if [[ "${current_mode}" != "${mode}" ]]; then
            if dry_run_action "Fix permissions on ${dir_path} from ${current_mode} to ${mode}"; then
                return 0
            fi
            chmod "${mode}" "${dir_path}"
            log_info "Fixed permissions on ${dir_path}: ${current_mode} -> ${mode}"
        fi
    else
        if dry_run_action "Create directory ${dir_path} with mode ${mode}"; then
            return 0
        fi
        mkdir -p "${dir_path}"
        chmod "${mode}" "${dir_path}"
        log_info "Created directory: ${dir_path} (mode ${mode})"
    fi
}

# backup_file PATH — creates timestamped backup
backup_file() {
    local file_path="${1:?missing file path}"

    if [[ ! -f "${file_path}" ]]; then
        log_debug "File does not exist, no backup needed: ${file_path}"
        return 0
    fi

    if [[ "${SKIP_BACKUP}" -eq 1 ]]; then
        log_warn "Skipping backup for ${file_path} (--no-backup)"
        return 0
    fi

    local backup_path="${file_path}${BACKUP_SUFFIX}"
    if dry_run_action "Backup ${file_path} -> ${backup_path}"; then
        return 0
    fi
    cp -a "${file_path}" "${backup_path}"
    log_info "Backed up ${file_path} -> ${backup_path}"
}

# ---------------------------------------------------------------------------
# Command Implementations
# ---------------------------------------------------------------------------

# cmd_help — display usage information
cmd_help() {
    cat <<'EOF'
ssh-setup.sh v1.0.0 — SSH Configuration Setup & Hardening Helper

USAGE:
    ./ssh-setup.sh [COMMAND] [OPTIONS]

COMMANDS:
    all           Run full setup workflow (check, generate key, configure)
    check         Check SSH installation and configuration status
    generate      Generate a new SSH key pair (Ed25519 preferred)
    add-key       Add a public key to authorized_keys
    harden        Harden /etc/ssh/sshd_config (requires root)
    status        Show current SSH configuration summary
    help          Show this help message

OPTIONS:
    -e, --email <email>     Email label for generated key
    -f, --file <path>       Path to key file (default: ~/.ssh/id_ed25519)
    -k, --key <pubkey>      Public key string to add
    -p, --port <port>       SSH port for hardening (default: 22)
    -n, --no-backup         Skip backup before hardening (dangerous)
    -y, --yes               Skip interactive prompts (non-interactive mode)
    -d, --dry-run           Show what would be done without making changes
    -v, --verbose           Enable verbose/debug output
    -h, --help              Show this help message

EXAMPLES:
    # Full automated setup
    ./ssh-setup.sh all

    # Generate a key with email label
    ./ssh-setup.sh generate --email admin@example.com

    # Add a specific public key
    ./ssh-setup.sh add-key --key "ssh-ed25519 AAAA... user@host"

    # Harden SSH on a custom port (non-interactive)
    ./ssh-setup.sh harden --port 2222 --yes

    # Preview what would happen (no changes)
    ./ssh-setup.sh all --dry-run

    # Check current status
    ./ssh-setup.sh status

SECURITY NOTES:
    - Ed25519 keys are preferred for security and performance
    - Hardening requires root privileges
    - Always backup sshd_config before making changes
    - Test SSH connectivity before closing current session after hardening

EXIT CODES:
    0   Success
    1   General error
    2   Invalid arguments
    3   Missing prerequisites
    4   Permission denied
EOF
}

# cmd_check — verify SSH installation and prerequisites
cmd_check() {
    log_info "Checking SSH installation and prerequisites..."

    local has_issues=0

    # Check SSH client
    if require_cmd ssh; then
        local ssh_version
        ssh_version="$(ssh -V 2>&1 || true)"
        log_info "SSH client: ${ssh_version}"
    else
        log_error "SSH client not installed"
        has_issues=1
    fi

    # Check SSH keygen
    if require_cmd ssh-keygen; then
        log_info "ssh-keygen: available"
    else
        log_error "ssh-keygen not found"
        has_issues=1
    fi

    # Check SSH agent
    if require_cmd ssh-add; then
        log_info "ssh-add: available"
    else
        log_warn "ssh-add not found"
    fi

    # Check SSH directory
    if [[ -d "${SSH_DIR}" ]]; then
        local ssh_mode
        ssh_mode="$(stat -c '%a' "${SSH_DIR}" 2>/dev/null || stat -f '%Lp' "${SSH_DIR}" 2>/dev/null)"
        if [[ "${ssh_mode}" == "700" ]]; then
            log_info "SSH directory: ${SSH_DIR} (permissions: ${ssh_mode} — OK)"
        else
            log_warn "SSH directory: ${SSH_DIR} (permissions: ${ssh_mode} — should be 700)"
        fi
    else
        log_warn "SSH directory does not exist: ${SSH_DIR}"
    fi

    # Check for existing keys
    local key_count=0
    for key_type in ed25519 rsa ecdsa; do
        local key_path="${SSH_DIR}/id_${key_type}"
        if [[ -f "${key_path}" ]]; then
            log_info "Found existing key: ${key_path}"
            ((key_count++)) || true
        fi
    done

    if [[ ${key_count} -eq 0 ]]; then
        log_warn "No SSH keys found in ${SSH_DIR}"
    else
        log_info "Found ${key_count} existing SSH key pair(s)"
    fi

    # Check sshd (server)
    if require_cmd sshd; then
        local sshd_version
        sshd_version="$(sshd -V 2>&1 || sshd -? 2>&1 || true)"
        log_info "SSH server (sshd): installed"
    else
        log_warn "SSH server (sshd) not installed — client only"
    fi

    # Check sshd_config
    if [[ -f "${DEFAULT_SSHD_CONFIG}" ]]; then
        log_info "sshd_config: ${DEFAULT_SSHD_CONFIG} (exists)"
    else
        log_warn "sshd_config not found at ${DEFAULT_SSHD_CONFIG}"
    fi

    if [[ ${has_issues} -ne 0 ]]; then
        log_warn "Some prerequisites are missing. Run './ssh-setup.sh all' to set up."
        return 1
    fi

    log_info "All SSH prerequisites are satisfied."
    return 0
}

# cmd_generate — create SSH key pair
cmd_generate() {
    local key_path="${KEY_FILE:-${DEFAULT_KEY_PATH}}"
    local key_type="${DEFAULT_KEY_TYPE}"
    local email_label="${USER_EMAIL:-"$(whoami)@$(hostname)"}"

    log_info "Generating SSH key pair (${key_type})..."
    log_info "Key path: ${key_path}"
    log_info "Label: ${email_label}"

    # Ensure .ssh directory exists
    ensure_dir "${SSH_DIR}" "700"

    # Check if key already exists
    if [[ -f "${key_path}" ]]; then
        log_warn "Key file already exists: ${key_path}"
        if ! confirm "Overwrite existing key?"; then
            log_info "Keeping existing key. Aborting key generation."
            return 0
        fi
        backup_file "${key_path}"
    fi

    if dry_run_action "Generate ${key_type} key at ${key_path} with label '${email_label}'"; then
        return 0
    fi

    require_cmd ssh-keygen || return 3

    # Generate Ed25519 key (preferred for security)
    local keygen_args=()
    keygen_args+=(-t "${key_type}")
    keygen_args+=(-f "${key_path}")
    keygen_args+=(-C "${email_label}")
    # No passphrase for non-interactive; prompt otherwise
    if [[ "${NON_INTERACTIVE}" -eq 1 ]]; then
        keygen_args+=(-N "")  # Empty passphrase
    fi

    ssh-keygen "${keygen_args[@]}" <<< "y" 2>&1 | while IFS= read -r line; do
        log_debug "ssh-keygen: ${line}"
    done

    # Verify key was created
    if [[ ! -f "${key_path}" ]]; then
        log_error "Key generation failed — key file not found: ${key_path}"
        return 1
    fi

    # Verify permissions
    chmod 600 "${key_path}"
    chmod 644 "${key_path}.pub"
    log_info "Private key permissions: 600 (owner read/write only)"
    log_info "Public key permissions: 644 (owner read/write, others read)"

    # Display public key
    log_info "Public key content:"
    printf "${BOLD}%s${NC}\n" "$(cat "${key_path}.pub")"

    log_info "Key pair generated successfully."
    return 0
}

# cmd_add_key — add public key to authorized_keys
cmd_add_key() {
    local auth_keys="${SSH_DIR}/authorized_keys"

    log_info "Managing authorized_keys..."

    # Ensure .ssh directory
    ensure_dir "${SSH_DIR}" "700"

    # If a specific key was provided via --key
    if [[ -n "${PUBLIC_KEY}" ]]; then
        _add_specific_key "${PUBLIC_KEY}" "${auth_keys}"
        return $?
    fi

    # Otherwise, offer to add the local public key
    local local_pub=""
    for key_type in ed25519 rsa ecdsa; do
        local key_path="${SSH_DIR}/id_${key_type}"
        if [[ -f "${key_path}.pub" ]]; then
            local_pub="${key_path}.pub"
            log_info "Found local public key: ${local_pub}"
            break
        fi
    done

    if [[ -z "${local_pub}" ]]; then
        log_error "No local public key found. Generate one first: ./ssh-setup.sh generate"
        return 1
    fi

    _add_specific_key "$(cat "${local_pub}")" "${auth_keys}"
    return $?
}

# _add_specific_key "key_content" auth_keys_path
_add_specific_key() {
    local key_content="${1:?missing key content}"
    local auth_keys="${2:?missing authorized_keys path}"

    # Ensure authorized_keys file exists with correct permissions
    if [[ ! -f "${auth_keys}" ]]; then
        if dry_run_action "Create ${auth_keys}"; then
            return 0
        fi
        touch "${auth_keys}"
        chmod 600 "${auth_keys}"
        log_info "Created ${auth_keys} (mode 600)"
    else
        local ak_mode
        ak_mode="$(stat -c '%a' "${auth_keys}" 2>/dev/null || stat -f '%Lp' "${auth_keys}" 2>/dev/null)"
        if [[ "${ak_mode}" != "600" ]]; then
            if dry_run_action "Fix authorized_keys permissions from ${ak_mode} to 600"; then
                return 0
            fi
            chmod 600 "${auth_keys}"
            log_info "Fixed authorized_keys permissions: ${ak_mode} -> 600"
        fi
    fi

    # Check for duplicate (compare the key type + base64 portion)
    local key_fingerprint
    key_fingerprint="$(echo "${key_content}" | awk '{print $1, $2}')"

    if grep -qF "${key_fingerprint}" "${auth_keys}" 2>/dev/null; then
        log_info "Key already present in ${auth_keys} — skipping (idempotent)"
        return 0
    fi

    if dry_run_action "Add public key to ${auth_keys}"; then
        return 0
    fi

    # Append key with a comment line
    printf '\n# Added by %s on %s\n%s\n' \
        "${SCRIPT_NAME}" \
        "$(date +"${DATE_FMT}")" \
        "${key_content}" >> "${auth_keys}"

    log_info "Public key added to ${auth_keys}"
    return 0
}

# cmd_harden — harden sshd_config
cmd_harden() {
    log_info "SSH Hardening Mode"

    # Require root
    if ! is_root; then
        log_error "Hardening requires root privileges. Run with: sudo ${SCRIPT_NAME} harden"
        return 4
    fi

    require_cmd sshd || return 3

    local sshd_config="${DEFAULT_SSHD_CONFIG}"

    if [[ ! -f "${sshd_config}" ]]; then
        log_error "sshd_config not found at ${sshd_config}"
        return 1
    fi

    # Backup
    backup_file "${sshd_config}"

    log_info "Applying SSH hardening settings to ${sshd_config}..."

    # Apply custom port first
    local current_port
    current_port="$(grep -i "^[[:space:]]*Port[[:space:]]" "${sshd_config}" 2>/dev/null | tail -1 | awk '{print $2}' || true)"

    if dry_run_action "Set Port ${SSH_PORT}"; then
        log_info "[DRY RUN] Would set Port to ${SSH_PORT}"
    elif [[ "${current_port}" == "${SSH_PORT}" ]]; then
        log_debug "Port already set to ${SSH_PORT}"
    elif [[ -n "${current_port}" ]]; then
        sed -i "s/^[[:space:]]*Port[[:space:]].*/Port ${SSH_PORT}/" "${sshd_config}"
        log_info "Updated: Port ${current_port} -> ${SSH_PORT}"
    else
        # No Port directive exists — add one explicitly
        # Try to uncomment "#Port 22" first, otherwise append
        if grep -qi "^[#[:space:]]*Port[[:space:]]\?22" "${sshd_config}" 2>/dev/null; then
            sed -i '0,/^[#[:space:]]*[Pp]ort/{s/^[#[:space:]]*[Pp]ort[[:space:]]\?[0-9]*/Port '"${SSH_PORT}"'/}' "${sshd_config}"
            log_info "Uncommented and set: Port ${SSH_PORT}"
        else
            # Insert Port directive near the top, after the first block of comments/defaults
            local insert_line
            insert_line="$(grep -n "^[^#]" "${sshd_config}" | head -1 | cut -d: -f1 || echo "1")"
            if [[ -n "${insert_line}" ]]; then
                sed -i "${insert_line}i\\
Port ${SSH_PORT}" "${sshd_config}"
            else
                printf '\n# SSH Port\nPort %s\n' "${SSH_PORT}" >> "${sshd_config}"
            fi
            log_info "Added: Port ${SSH_PORT}"
        fi
    fi

    # Hardening directives: key=value pairs
    # Format: directive|value|comment
    declare -a hardening_rules=(
        "PermitRootLogin|no|Disable root login via SSH"
        "PasswordAuthentication|yes|Enable password authentication"
        "KbdInteractiveAuthentication|yes|Enable keyboard-interactive authentication"
        "PubkeyAuthentication|yes|Enable public key authentication"
        "ChallengeResponseAuthentication|no|Disable challenge-response auth"
        "UsePAM|yes|Enable PAM for account management"
        "X11Forwarding|no|Disable X11 forwarding"
        "AllowTcpForwarding|no|Disable TCP forwarding"
        "PermitTunnel|no|Disable tunneling"
        "MaxAuthTries|3|Limit authentication attempts"
        "MaxSessions|5|Limit concurrent sessions"
        "LoginGraceTime|30|Reduce login grace time (seconds)"
        "ClientAliveInterval|300|Send keepalive every 5 minutes"
        "ClientAliveCountMax|2|Disconnect after 2 missed keepalives"
        "AllowAgentForwarding|no|Disable agent forwarding"
        "PermitEmptyPasswords|no|Disallow empty passwords"
        "IgnoreRhosts|yes|Ignore .rhosts files"
        "HostbasedAuthentication|no|Disable host-based auth"
        "PrintMotd|no|Use PAM motd instead"
    )

    local changes_made=0

    for rule in "${hardening_rules[@]}"; do
        IFS='|' read -r directive value comment <<< "${rule}"

        if dry_run_action "Set ${directive} ${value}  # ${comment}"; then
            ((changes_made++)) || true
            continue
        fi

        # Check if directive already exists with correct value
        local current_val
        current_val="$(grep -i "^[[:space:]]*${directive}[[:space:]]" "${sshd_config}" 2>/dev/null | tail -1 | awk '{print $2}' || true)"

        if [[ "${current_val}" == "${value}" ]]; then
            log_debug "Already set: ${directive} ${value}"
            continue
        fi

        if [[ -n "${current_val}" ]]; then
            # Directive exists with different value — update in place
            sed -i "s/^[[:space:]]*${directive}[[:space:]].*/${directive} ${value}/" "${sshd_config}"
            log_info "Updated: ${directive} ${current_val} -> ${value}  # ${comment}"
        else
            # Directive doesn't exist — append
            printf '\n# %s\n%s %s\n' "${comment}" "${directive}" "${value}" >> "${sshd_config}"
            log_info "Added: ${directive} ${value}  # ${comment}"
        fi

        ((changes_made++)) || true
    done

    # Protocol 2 only (for older OpenSSH — modern versions ignore this)
    local openssh_version
    openssh_version="$(sshd -V 2>&1 | grep -oP '[0-9]+\.[0-9]+' | head -1 || echo "unknown")"
    log_debug "OpenSSH version: ${openssh_version}"

    if [[ ${changes_made} -eq 0 ]]; then
        log_info "All hardening settings already applied. No changes needed."
    else
        log_info "Applied ${changes_made} hardening change(s)."
        log_warn "IMPORTANT: Before closing your current session, verify SSH connectivity:"
        log_warn "  ssh -p ${SSH_PORT} localhost"
        log_warn "  ssh -p ${SSH_PORT} user@your-server"
    fi

    # Ensure AuthenticationMethods doesn't override password auth
    # This directive can force specific auth methods and bypass PasswordAuthentication
    if grep -qi "^[[:space:]]*AuthenticationMethods[[:space:]]" "${sshd_config}" 2>/dev/null; then
        if dry_run_action "Comment out AuthenticationMethods directive to allow password auth"; then
            log_info "[DRY RUN] Would comment out AuthenticationMethods"
        else
            # Comment out any AuthenticationMethods lines
            sed -i 's/^[[:space:]]*AuthenticationMethods[[:space:]].*/# &  # Commented by ssh-setup.sh to allow password auth/' "${sshd_config}"
            log_info "Commented out AuthenticationMethods directive (can override password auth)"
        fi
    fi

    # Ensure no Match block overrides disable password auth
    # Add explicit Match all block to guarantee password auth is enabled globally
    if dry_run_action "Add 'Match all' block to ensure password auth is enabled"; then
        log_info "[DRY RUN] Would add 'Match all' block"
    else
        # Check if there's already a Match all block
        if ! grep -qi "^[[:space:]]*Match[[:space:]]\+all" "${sshd_config}" 2>/dev/null; then
            printf '\n# Ensure password authentication is enabled for all users\nMatch all\n  PasswordAuthentication yes\n  KbdInteractiveAuthentication yes\n' >> "${sshd_config}"
            log_info "Added 'Match all' block to guarantee password authentication"
        else
            log_debug "Match all block already present"
        fi
    fi

    # Validate config syntax
    if dry_run_action "Validate sshd_config syntax"; then
        return 0
    fi

    log_info "Validating sshd_config syntax..."
    if sshd -t 2>&1; then
        log_info "sshd_config syntax is valid."
    else
        log_error "sshd_config has syntax errors! Restoring backup..."
        if [[ -f "${sshd_config}${BACKUP_SUFFIX}" ]]; then
            cp -a "${sshd_config}${BACKUP_SUFFIX}" "${sshd_config}"
            log_info "Restored from backup."
        fi
        return 1
    fi

    # Offer to restart SSH service
    if [[ "${NON_INTERACTIVE}" -eq 0 ]]; then
        if confirm "Restart SSH service now? (This may disconnect you)"; then
            log_action "Restarting SSH service..."
            if command -v systemctl &>/dev/null; then
                systemctl restart sshd || systemctl restart ssh || {
                    log_error "Failed to restart SSH service"
                    return 1
                }
            else
                service sshd restart || service ssh restart || {
                    log_error "Failed to restart SSH service"
                    return 1
                }
            fi
            log_info "SSH service restarted successfully."
        else
            log_info "SSH service NOT restarted. Remember to restart it manually after testing."
            log_info "  sudo systemctl restart sshd   # or: sudo service ssh restart"
        fi
    fi

    return 0
}

# cmd_status — show current SSH configuration summary
cmd_status() {
    log_info "SSH Configuration Status"
    printf "${BOLD}========================================${NC}\n" >&2

    # SSH client
    if command -v ssh &>/dev/null; then
        local ssh_ver
        ssh_ver="$(ssh -V 2>&1 | head -1 || true)"
        printf "${GREEN}[✓]${NC} SSH Client: %s\n" "${ssh_ver}" >&2
    else
        printf "${RED}[✗]${NC} SSH Client: NOT INSTALLED\n" >&2
    fi

    # SSH server
    if command -v sshd &>/dev/null; then
        printf "${GREEN}[✓]${NC} SSH Server: installed\n" >&2
    else
        printf "${YELLOW}[!]${NC} SSH Server: not installed\n" >&2
    fi

    # SSH directory
    if [[ -d "${SSH_DIR}" ]]; then
        local ssh_mode
        ssh_mode="$(stat -c '%a' "${SSH_DIR}" 2>/dev/null || stat -f '%Lp' "${SSH_DIR}" 2>/dev/null)"
        if [[ "${ssh_mode}" == "700" ]]; then
            printf "${GREEN}[✓]${NC} ~/.ssh: exists (mode %s)\n" "${ssh_mode}" >&2
        else
            printf "${RED}[✗]${NC} ~/.ssh: exists but mode is %s (should be 700)\n" "${ssh_mode}" >&2
        fi
    else
        printf "${RED}[✗]${NC} ~/.ssh: does not exist\n" >&2
    fi

    # Keys
    printf "\n${BOLD}SSH Keys:${NC}\n" >&2
    local found_keys=0
    for key_type in ed25519 rsa ecdsa; do
        local key_path="${SSH_DIR}/id_${key_type}"
        if [[ -f "${key_path}" ]]; then
            local fp
            fp="$(ssh-keygen -lf "${key_path}" 2>/dev/null || echo "unable to read fingerprint")"
            printf "  ${GREEN}[✓]${NC} %s (%s-bit)\n" "${key_type}" "${fp}" >&2
            ((found_keys++)) || true
        fi
    done
    if [[ ${found_keys} -eq 0 ]]; then
        printf "  ${YELLOW}[!]${NC} No SSH keys found\n" >&2
    fi

    # authorized_keys
    local ak="${SSH_DIR}/authorized_keys"
    if [[ -f "${ak}" ]]; then
        local ak_count
        ak_count="$(grep -c '^ssh-' "${ak}" 2>/dev/null || echo "0")"
        printf "\n${GREEN}[✓]${NC} authorized_keys: %s key(s)\n" "${ak_count}" >&2
    else
        printf "\n${YELLOW}[!]${NC} authorized_keys: not found\n" >&2
    fi

    # sshd_config summary — show EFFECTIVE config (what sshd actually uses)
    if command -v sshd &>/dev/null; then
        printf "\n${BOLD}sshd_config Effective Settings (active):${NC}\n" >&2
        local directives=("port" "permitrootlogin" "passwordauthentication" "kbdinteractiveauthentication" "pubkeyauthentication" "maxauthtries" "x11forwarding" "authenticationmethods" "challengeresponseauthentication" "allowtcpforwarding" "permittunnel" "allowagentforwarding" "permitemptypasswords" "ignorerhosts" "hostbasedauthentication" "printmotd" "usepam" "clientaliveinterval" "clientalivecountmax" "logingracetime" "maxsessions")
        local display_names=("Port" "PermitRootLogin" "PasswordAuthentication" "KbdInteractiveAuthentication" "PubkeyAuthentication" "MaxAuthTries" "X11Forwarding" "AuthenticationMethods" "ChallengeResponseAuthentication" "AllowTcpForwarding" "PermitTunnel" "AllowAgentForwarding" "PermitEmptyPasswords" "IgnoreRhosts" "HostbasedAuthentication" "PrintMotd" "UsePAM" "ClientAliveInterval" "ClientAliveCountMax" "LoginGraceTime" "MaxSessions")

        # sshd -T requires root for full output; fall back to file grep if not root
        if is_root; then
            local effective_config
            effective_config="$(sshd -T 2>/dev/null || true)"
            if [[ -n "${effective_config}" ]]; then
                for i in "${!directives[@]}"; do
                    local key="${directives[$i]}"
                    local label="${display_names[$i]}"
                    local val
                    val="$(echo "${effective_config}" | grep -i "^${key} " | awk '{print $2}' || true)"
                    if [[ -z "${val}" ]]; then
                        val="not set (default)"
                    fi
                    printf "  %-40s %s\n" "${label}" "${val}" >&2
                done
            else
                printf "  ${RED}[✗]${NC} Could not read effective config (try sudo)\n" >&2
            fi
        else
            # Non-root: grep the config file for uncommented active directives
            printf "  ${YELLOW}[!]${NC} Showing file content (run with sudo for effective config)\n" >&2
            for i in "${!directives[@]}"; do
                local key="${directives[$i]}"
                local label="${display_names[$i]}"
                local val
                val="$(grep -i "^[[:space:]]*${key}[[:space:]]" "${DEFAULT_SSHD_CONFIG}" 2>/dev/null | tail -1 | awk '{print $2}' || echo "not set")"
                printf "  %-40s %s\n" "${label}" "${val}" >&2
            done
        fi
    else
        printf "\n${YELLOW}[!]${NC} sshd: not installed\n" >&2
    fi

    printf "\n${BOLD}========================================${NC}\n" >&2
    log_info "Status check complete."
    return 0
}

# cmd_all — run full setup workflow
cmd_all() {
    log_info "Running full SSH setup workflow..."
    log_info "========================================"

    # Step 1: Check
    log_info "Step 1/4: Checking prerequisites..."
    cmd_check || true  # Continue even if some checks fail

    # Step 2: Generate key if none exists
    local has_ed25519=0
    [[ -f "${SSH_DIR}/id_ed25519" ]] && has_ed25519=1

    if [[ ${has_ed25519} -eq 0 ]]; then
        log_info "Step 2/4: Generating SSH key..."
        cmd_generate
    else
        log_info "Step 2/4: Ed25519 key already exists — skipping generation."
    fi

    # Step 3: Configure authorized_keys
    log_info "Step 3/4: Configuring authorized_keys..."
    cmd_add_key

    # Step 4: Show status
    log_info "Step 4/4: Final status..."
    cmd_status

    log_info "========================================"
    log_info "SSH setup complete!"
    log_info "Your public key is at: ${SSH_DIR}/id_ed25519.pub"
    log_info "Copy it to remote servers with: ssh-copy-id user@host"

    return 0
}

# ---------------------------------------------------------------------------
# Argument Parsing
# ---------------------------------------------------------------------------

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            all|check|generate|add-key|harden|status|help)
                COMMAND="$1"
                shift
                ;;
            -e|--email)
                [[ $# -ge 2 ]] || { log_error "--email requires a value"; exit 2; }
                USER_EMAIL="$2"
                shift 2
                ;;
            -f|--file)
                [[ $# -ge 2 ]] || { log_error "--file requires a value"; exit 2; }
                KEY_FILE="$2"
                shift 2
                ;;
            -k|--key)
                [[ $# -ge 2 ]] || { log_error "--key requires a value"; exit 2; }
                PUBLIC_KEY="$2"
                shift 2
                ;;
            -p|--port)
                [[ $# -ge 2 ]] || { log_error "--port requires a value"; exit 2; }
                if [[ ! "$2" =~ ^[0-9]+$ ]] || [[ "$2" -lt 1 ]] || [[ "$2" -gt 65535 ]]; then
                    log_error "--port must be a number between 1 and 65535"
                    exit 2
                fi
                SSH_PORT="$2"
                shift 2
                ;;
            -n|--no-backup)
                SKIP_BACKUP=1
                shift
                ;;
            -y|--yes)
                NON_INTERACTIVE=1
                shift
                ;;
            -d|--dry-run)
                DRY_RUN=1
                shift
                ;;
            -v|--verbose)
                VERBOSE=1
                shift
                ;;
            -h|--help)
                cmd_help
                exit 0
                ;;
            --)
                shift
                break
                ;;
            -*)
                log_error "Unknown option: $1"
                cmd_help
                exit 2
                ;;
            *)
                if [[ -z "${COMMAND}" ]]; then
                    COMMAND="$1"
                else
                    log_error "Unexpected argument: $1"
                    exit 2
                fi
                shift
                ;;
        esac
    done

    # Default command
    if [[ -z "${COMMAND}" ]]; then
        COMMAND="help"
    fi
}

# ---------------------------------------------------------------------------
# Main Entry Point
# ---------------------------------------------------------------------------

main() {
    parse_args "$@"

    log_debug "Command: ${COMMAND}"
    log_debug "Dry run: ${DRY_RUN}"
    log_debug "Non-interactive: ${NON_INTERACTIVE}"
    log_debug "Verbose: ${VERBOSE}"

    # Route to command handler
    case "${COMMAND}" in
        all)
            cmd_all
            ;;
        check)
            cmd_check
            ;;
        generate)
            cmd_generate
            ;;
        add-key)
            cmd_add_key
            ;;
        harden)
            cmd_harden
            ;;
        status)
            cmd_status
            ;;
        help|"")
            cmd_help
            ;;
        *)
            log_error "Unknown command: ${COMMAND}"
            cmd_help
            exit 2
            ;;
    esac
}

main "$@"