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

usage() {
    cat >&2 <<'MSG'
Usage:
  niceos-k8s-set-root-key 'ssh-ed25519 AAAA... comment'
  niceos-k8s-set-root-key - < /path/to/id_ed25519.pub

This command installs /root/.ssh/authorized_keys and keeps temporary
password SSH enabled until a successful root publickey SSH login is observed.

After running it, open a second terminal and test:
  ssh -o PreferredAuthentications=publickey -o PasswordAuthentication=no root@NODE_IP
MSG
}

if [ "${1:-}" = "-h" ] || [ "${1:-}" = "--help" ]; then
    usage
    exit 0
fi

key=""

if [ "${1:-}" = "-" ]; then
    key="$(cat)"
elif [ $# -ge 1 ]; then
    key="$*"
else
    usage
    exit 2
fi

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

if ! grep -Eq '(ssh-ed25519|ssh-rsa|rsa-sha2-256|rsa-sha2-512|ecdsa-sha2-|sk-ssh-ed25519|sk-ecdsa-sha2-)' "${tmp}"; then
    rm -f "${tmp}"
    echo "ERROR: input does not look like an SSH public key" >&2
    exit 1
fi

if ! ssh-keygen -l -f "${tmp}" >/dev/null 2>&1; then
    rm -f "${tmp}"
    echo "ERROR: ssh-keygen could not parse this public key" >&2
    exit 1
fi

install -d -m 0700 /root/.ssh
install -m 0600 "${tmp}" /root/.ssh/authorized_keys
rm -f "${tmp}"

chown -R root:root /root/.ssh
chmod 0700 /root/.ssh
chmod 0600 /root/.ssh/authorized_keys

if [ -x /usr/libexec/niceos-k8s-ssh-enrollment ]; then
    /usr/libexec/niceos-k8s-ssh-enrollment reconcile
fi

cat <<'MSG'
OK: root SSH key installed and validated syntactically.
INFO: temporary root password SSH is still enabled until the key is proven.

Open a second terminal and test key-only login:
  ssh -o PreferredAuthentications=publickey -o PasswordAuthentication=no root@NODE_IP

After successful publickey login, NiceOS will automatically remove:
  /etc/ssh/sshd_config.d/20-niceos-k8s-firstboot-root-password.conf

Check status:
  niceos-k8s-ssh-enrollment-status
MSG