#!/usr/bin/env bash
set -euo pipefail

state_dir="/var/lib/niceos/k8s-ssh-enrollment"
log_file="${state_dir}/enrollment.log"

tmp_conf="/etc/ssh/sshd_config.d/21-niceos-k8s-firstboot-root-password.conf"
root_ssh_dir="/root/.ssh"
root_auth_keys="/root/.ssh/authorized_keys"

key_valid_marker="${state_dir}/key.valid"
key_verified_marker="${state_dir}/key.verified"
done_marker="${state_dir}/done"
key_installed_since_file="${state_dir}/key-installed-since"

install -d -m 0755 "${state_dir}"

log() {
    printf '%s %s\n' "$(date -Is)" "$*" | tee -a "${log_file}" >&2
}

reload_sshd_if_running() {
    if command -v sshd >/dev/null 2>&1; then
        if ! sshd -t >/dev/null 2>&1; then
            log "[WARN] sshd -t failed; sshd was not reloaded"
            return 1
        fi
    fi

    if command -v systemctl >/dev/null 2>&1; then
        if systemctl is-active --quiet sshd.service; then
            systemctl reload sshd.service >/dev/null 2>&1 || \
            systemctl restart sshd.service >/dev/null 2>&1 || :
        fi
    fi

    return 0
}

extract_key_from_line() {
    # Prints "type base64" if a public key token is found in an authorized_keys line.
    # Supports lines with simple authorized_keys options before the key.
    awk '
    BEGIN {
        key_re="^(ssh-ed25519|ssh-rsa|rsa-sha2-256|rsa-sha2-512|ecdsa-sha2-|sk-ssh-ed25519|sk-ecdsa-sha2-)"
    }
    /^[[:space:]]*$/ { exit 1 }
    /^[[:space:]]*#/ { exit 1 }
    {
        for (i = 1; i <= NF; i++) {
            if ($i ~ key_re && (i + 1) <= NF) {
                print $i, $(i + 1)
                exit 0
            }
        }
        exit 1
    }'
}

line_has_valid_public_key() {
    local line="$1"
    local tmp
    local parsed

    parsed="$(printf '%s\n' "${line}" | extract_key_from_line || true)"
    [ -n "${parsed}" ] || return 1

    tmp="$(mktemp)"
    printf '%s\n' "${parsed}" > "${tmp}"

    if ssh-keygen -l -f "${tmp}" >/dev/null 2>&1; then
        rm -f "${tmp}"
        return 0
    fi

    rm -f "${tmp}"
    return 1
}

authorized_keys_has_valid_key() {
    local line
    local valid_count=0

    [ -s "${root_auth_keys}" ] || return 1

    while IFS= read -r line || [ -n "${line}" ]; do
        if line_has_valid_public_key "${line}"; then
            valid_count=$((valid_count + 1))
        fi
    done < "${root_auth_keys}"

    [ "${valid_count}" -gt 0 ]
}

fix_root_ssh_permissions() {
    if [ -d "${root_ssh_dir}" ]; then
        chown root:root "${root_ssh_dir}" || :
        chmod 0700 "${root_ssh_dir}" || :
    fi

    if [ -e "${root_auth_keys}" ]; then
        chown root:root "${root_auth_keys}" || :
        chmod 0600 "${root_auth_keys}" || :
    fi
}

write_temporary_password_dropin() {
    install -d -m 0755 /etc/ssh/sshd_config.d

    cat > "${tmp_conf}" <<'CONF'
# NiceOS K8s Node temporary SSH enrollment mode.
# This file is generated automatically and removed only after a successful
# root publickey SSH login is observed.

PermitRootLogin yes
PubkeyAuthentication yes

# "publickey password" means either publickey OR password.
# Do not use "publickey,password" here.
AuthenticationMethods publickey password

PasswordAuthentication yes
KbdInteractiveAuthentication no
PermitEmptyPasswords no

UsePAM yes
StrictModes yes
AuthorizedKeysFile .ssh/authorized_keys
CONF

    chmod 0644 "${tmp_conf}"

    log "[OK] temporary root password SSH enrollment drop-in is present: ${tmp_conf}"
    reload_sshd_if_running || :
}

mark_key_valid_unverified() {
    local now

    fix_root_ssh_permissions

    touch "${key_valid_marker}"

    if [ ! -s "${key_installed_since_file}" ]; then
        now="$(date +%s)"
        printf '%s\n' "${now}" > "${key_installed_since_file}"
    fi

    log "[OK] valid root authorized_keys detected; waiting for successful publickey login"
}

clear_key_valid_state() {
    rm -f "${key_valid_marker}"
    rm -f "${key_installed_since_file}"
}

secure_after_verified_publickey_login() {
    rm -f "${tmp_conf}"
    touch "${key_verified_marker}"
    touch "${done_marker}"

    log "[OK] successful root publickey login detected; temporary password SSH drop-in removed"
    reload_sshd_if_running || :
}

journal_has_successful_root_publickey_login() {
    local since
    local since_arg

    [ -s "${key_installed_since_file}" ] || return 1
    since="$(cat "${key_installed_since_file}" 2>/dev/null || true)"
    [ -n "${since}" ] || return 1

    since_arg="@${since}"

    if ! command -v journalctl >/dev/null 2>&1; then
        return 1
    fi

    # OpenSSH normally logs:
    # Accepted publickey for root from X.X.X.X port N ssh2: ...
    journalctl -b \
        -u sshd.service \
        --since "${since_arg}" \
        -o cat \
        --no-pager 2>/dev/null \
        | grep -Eq 'Accepted publickey for root[[:space:]]'
}

reconcile() {
    if [ -e "${done_marker}" ]; then
        rm -f "${tmp_conf}"
        log "[INFO] enrollment already completed; password SSH will not be re-enabled"
        reload_sshd_if_running || :
        return 0
    fi

    if authorized_keys_has_valid_key; then
        mark_key_valid_unverified

        if journal_has_successful_root_publickey_login; then
            secure_after_verified_publickey_login
            return 0
        fi

        # Important: keep password SSH enabled until the key is proven by login.
        write_temporary_password_dropin
        return 0
    fi

    clear_key_valid_state
    log "[WARN] root authorized_keys is missing, empty or invalid; keeping password enrollment mode"
    write_temporary_password_dropin
}

verify() {
    if [ -e "${done_marker}" ]; then
        rm -f "${tmp_conf}"
        return 0
    fi

    if ! authorized_keys_has_valid_key; then
        log "[WARN] cannot verify publickey login: authorized_keys is missing or invalid"
        write_temporary_password_dropin
        return 0
    fi

    mark_key_valid_unverified

    if journal_has_successful_root_publickey_login; then
        secure_after_verified_publickey_login
    else
        log "[INFO] valid key exists, but no successful root publickey login observed yet"
        write_temporary_password_dropin
    fi
}

force_secure() {
    if ! authorized_keys_has_valid_key; then
        log "[ERROR] refusing force-secure: /root/.ssh/authorized_keys has no valid public key"
        exit 1
    fi

    rm -f "${tmp_conf}"
    touch "${key_valid_marker}"
    touch "${key_verified_marker}"
    touch "${done_marker}"

    log "[OK] forced secure mode; temporary password SSH drop-in removed"
    reload_sshd_if_running || :
}

status() {
    echo "state_dir=${state_dir}"
    echo "done=$([ -e "${done_marker}" ] && echo yes || echo no)"
    echo "key_valid_marker=$([ -e "${key_valid_marker}" ] && echo yes || echo no)"
    echo "key_verified_marker=$([ -e "${key_verified_marker}" ] && echo yes || echo no)"
    echo "temporary_dropin=$([ -e "${tmp_conf}" ] && echo yes || echo no)"
    echo "authorized_keys_exists=$([ -e "${root_auth_keys}" ] && echo yes || echo no)"
    echo "authorized_keys_nonempty=$([ -s "${root_auth_keys}" ] && echo yes || echo no)"

    if authorized_keys_has_valid_key; then
        echo "authorized_keys_valid=yes"
    else
        echo "authorized_keys_valid=no"
    fi

    if journal_has_successful_root_publickey_login; then
        echo "successful_publickey_login_seen=yes"
    else
        echo "successful_publickey_login_seen=no"
    fi

    if [ -s "${key_installed_since_file}" ]; then
        echo "key_installed_since_epoch=$(cat "${key_installed_since_file}")"
    else
        echo "key_installed_since_epoch="
    fi
}

case "${1:-reconcile}" in
    reconcile)
        reconcile
        ;;

    verify)
        verify
        ;;

    force-secure)
        force_secure
        ;;

    status)
        status
        ;;

    *)
        echo "Usage: $0 {reconcile|verify|force-secure|status}" >&2
        exit 2
        ;;
esac