#!/bin/bash

# LXS - Root SSH password login
# Description: Enable or disable root login over SSH with a password
# Author: LXS
# Date: 2025

# Load LXS common library (colors, separator, run_spinner, loggers, helpers)
LXS_RAW_BASE="${LXS_RAW_BASE:-https://git.hyko.cx/hykocx/lxs/raw/branch/main}"
_lib=$(curl -fsSL "${LXS_RAW_BASE}/lib/common.sh") || { echo "Failed to fetch lib/common.sh" >&2; exit 1; }
eval "$_lib"
unset _lib
export LXS_LOG_FILE="/tmp/lxs_root_ssh_login.log"

require_root "$0" "$@"

set -u

# Drop-in path. Numeric prefix `00-` makes it win over distro defaults — sshd
# applies the first match per option across the included files.
DROPIN_FILE="/etc/ssh/sshd_config.d/00-lxs-root-login.conf"
SSH_SERVICE="ssh"
command -v systemctl >/dev/null 2>&1 && systemctl list-unit-files 2>/dev/null | grep -q '^sshd\.service' && SSH_SERVICE="sshd"

# ═══════════════════════════════════════════════════════════════════════════
# Arguments
# ═══════════════════════════════════════════════════════════════════════════

ACTION=""

for arg in "$@"; do
    case "$arg" in
        --enable|enable)     ACTION="enable" ;;
        --disable|disable)   ACTION="disable" ;;
        --status|status)     ACTION="status" ;;
        -h|--help)
            cat <<EOF
Usage: root-ssh-login.sh [action]

Actions:
  --enable     Allow root to log in over SSH with a password
  --disable    Disallow root password login (restore SSH defaults)
  --status     Show the current effective settings
  -h, --help   Show this help

With no action, an interactive menu is shown.

Note: the change is written to ${DROPIN_FILE}
and applied by reloading the SSH service after a successful 'sshd -t'.
EOF
            exit 0
            ;;
        *)
            echo -e "${RED}Unknown option: $arg${NC}" >&2
            exit 1
            ;;
    esac
done

# ═══════════════════════════════════════════════════════════════════════════
# Pre-checks
# ═══════════════════════════════════════════════════════════════════════════

if [ ! -d /etc/ssh ] || ! command -v sshd >/dev/null 2>&1; then
    err "OpenSSH server is not installed."
    exit 1
fi

# ═══════════════════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════════════════

current_setting() {
    local key=$1
    sshd -T 2>/dev/null | awk -v k="${key,,}" 'tolower($1)==k {print $2; exit}'
}

show_status() {
    local permit pwauth
    permit=$(current_setting PermitRootLogin)
    pwauth=$(current_setting PasswordAuthentication)
    echo -e "${WHITE}${BOLD}Current SSH login settings${NC}"
    show_separator
    echo -e "  PermitRootLogin:        ${BOLD}${permit:-unknown}${NC}"
    echo -e "  PasswordAuthentication: ${BOLD}${pwauth:-unknown}${NC}"
    if [ -f "$DROPIN_FILE" ]; then
        echo -e "  Drop-in:                ${GRAY}${DROPIN_FILE}${NC}"
    else
        echo -e "  Drop-in:                ${GRAY}(none — distro defaults)${NC}"
    fi
    show_separator
    if [ "${permit:-}" = "yes" ] && [ "${pwauth:-}" = "yes" ]; then
        warn "Root password login is currently ENABLED."
    else
        ok "Root password login is currently DISABLED."
    fi
}

reload_sshd() {
    if ! sshd -t 2>>"$LXS_LOG_FILE"; then
        err "sshd config test failed — see ${LXS_LOG_FILE}. Reverting."
        return 1
    fi
    if systemctl reload "$SSH_SERVICE" 2>>"$LXS_LOG_FILE"; then
        ok "${SSH_SERVICE} reloaded"
        return 0
    fi
    # Fall back to restart (some minimal images ship without reload support)
    if systemctl restart "$SSH_SERVICE" 2>>"$LXS_LOG_FILE"; then
        ok "${SSH_SERVICE} restarted"
        return 0
    fi
    err "Failed to reload/restart ${SSH_SERVICE} — see ${LXS_LOG_FILE}"
    return 1
}

enable_root_login() {
    # Warn if root has no password — enabling password auth would be useless.
    local pw_status
    pw_status=$(passwd -S root 2>/dev/null | awk '{print $2}')
    case "$pw_status" in
        L|NP)
            warn "Root account has no usable password (status: ${pw_status})."
            warn "Run 'lxs tool root-password' first, otherwise login will still fail."
            ;;
    esac

    cat > "${DROPIN_FILE}.tmp" <<'EOF'
# Managed by LXS (lxs tool root-ssh-login). Remove this file to revert.
PermitRootLogin yes
PasswordAuthentication yes
EOF
    chmod 644 "${DROPIN_FILE}.tmp"
    mv "${DROPIN_FILE}.tmp" "$DROPIN_FILE"

    if ! reload_sshd; then
        rm -f "$DROPIN_FILE"
        reload_sshd >/dev/null 2>&1 || true
        return 1
    fi
    ok "Root password login over SSH is now ENABLED"
    warn "This weakens server security. Disable it again when no longer needed:"
    echo -e "  ${GRAY}lxs tool root-ssh-login --disable${NC}"
}

disable_root_login() {
    if [ ! -f "$DROPIN_FILE" ]; then
        info "Drop-in not present — root password login already follows SSH defaults."
    else
        local backup="${DROPIN_FILE}.bak.$$"
        mv "$DROPIN_FILE" "$backup"
        if ! reload_sshd; then
            mv "$backup" "$DROPIN_FILE"
            reload_sshd >/dev/null 2>&1 || true
            return 1
        fi
        rm -f "$backup"
    fi
    ok "Root password login over SSH is now DISABLED"
    echo ""
    show_status
}

# ═══════════════════════════════════════════════════════════════════════════
# Menu (when no action is given on the CLI)
# ═══════════════════════════════════════════════════════════════════════════

if [ -z "$ACTION" ]; then
    show_status
    echo ""
    echo -e "  ${GREEN}[1]${NC} Enable root SSH password login"
    echo -e "  ${YELLOW}[2]${NC} Disable root SSH password login"
    echo -e "  ${RED}[0]${NC} Cancel"
    echo ""
    echo -e -n "${BOLD}Choice [0-2]: ${NC}"
    read -r choice
    case "$choice" in
        1) ACTION="enable" ;;
        2) ACTION="disable" ;;
        0|"") info "Cancelled."; exit 0 ;;
        *) err "Invalid option."; exit 1 ;;
    esac
fi

case "$ACTION" in
    enable)  enable_root_login ;;
    disable) disable_root_login ;;
    status)  show_status ;;
esac