#!/usr/bin/env bash
# =============================================================================
# onx-ssh-key-add — Append an SSH public key to a user's authorized_keys.
#
# Input (stdin JSON):
#   {
#     "username": "onx_acme01",                  -- required
#     "key":      "ssh-rsa AAAA... user@host",   -- required (full single line)
#     "comment":  "my-laptop"                    -- optional, replaces inline comment
#   }
#
# Output (stdout JSON):
#   {
#     "username":    "onx_acme01",
#     "fingerprint": "SHA256:abc...",
#     "key_type":    "ssh-rsa",
#     "total_keys":  3,
#     "added":       true
#   }
#
# Idempotency:
#   If the same key (matched by fingerprint) is already present, returns
#   added=false instead of duplicating the line.
#
# Filesystem layout enforced:
#   /home/<user>/.ssh           mode 0700, owned by <user>:<user>
#   /home/<user>/.ssh/authorized_keys  mode 0600, owned by <user>:<user>
#
# Exit codes: 0=ok 1=invalid 2=preflight 3=exec 4=rolled-back 5=rollback-fail
#
# Deployed to: /usr/local/onoxsoft/bin/onx-ssh-key-add
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
# shellcheck source=_lib/common.sh
source "${SCRIPT_DIR}/_lib/common.sh"

require_root
require_cmd ssh-keygen

onx_json_input

USERNAME=$(onx_json_field "username")
KEY=$(onx_json_field      "key")
COMMENT=$(onx_json_field  "comment" "")

# ── Input validation ─────────────────────────────────────────────────────────
onx_validate_username "${USERNAME}"
[[ -z "${KEY}" ]] && onx_die 1 "key is required"

# Strip leading/trailing whitespace; collapse any embedded newlines (single-line key)
KEY="$(printf '%s' "${KEY}" | tr -d '\r' | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')"
case "${KEY}" in
    *$'\n'*) onx_die 1 "key must be a single line (got embedded newline)" ;;
esac

# Quick prefix sanity check before invoking ssh-keygen
case "${KEY}" in
    ssh-rsa\ *|ssh-ed25519\ *|ssh-ecdsa\ *|ecdsa-sha2-*\ *|sk-ssh-ed25519@*\ *|sk-ecdsa-sha2-*\ *) : ;;
    *) onx_die 1 "key has unrecognised type prefix (expected ssh-rsa / ssh-ed25519 / ecdsa-sha2-*)" ;;
esac

# Extract key_type early (first whitespace-separated token)
KEY_TYPE="${KEY%% *}"

# ── Validate via ssh-keygen + capture fingerprint ────────────────────────────
# We write to a temp file so we never expose the raw key on argv.
TMP_KEY="$(mktemp -t onx-ssh.XXXXXX)"
chmod 600 "${TMP_KEY}"
trap '[[ -n "${TMP_KEY:-}" ]] && rm -f "${TMP_KEY}" 2>/dev/null || true' EXIT

printf '%s\n' "${KEY}" > "${TMP_KEY}"

# ssh-keygen -l -f <file> emits "<bits> SHA256:<hash> <comment> (<type>)"
KEYGEN_OUT="$(ssh-keygen -l -f "${TMP_KEY}" 2>/dev/null || true)"
if [[ -z "${KEYGEN_OUT}" ]]; then
    onx_die 1 "ssh-keygen rejected the key (invalid format or unsupported type)"
fi

FINGERPRINT="$(printf '%s' "${KEYGEN_OUT}" | awk '{print $2}')"
[[ "${FINGERPRINT}" =~ ^SHA256: ]] || onx_die 1 "could not extract SHA256 fingerprint from ssh-keygen output"

# ── Preflight: account home must exist (provisioned user) ────────────────────
HOME_DIR="/home/${USERNAME}"
SSH_DIR="${HOME_DIR}/.ssh"
AK_FILE="${SSH_DIR}/authorized_keys"

if [[ "${MOCK_MODE}" != "1" ]]; then
    [[ -d "${HOME_DIR}" ]] || onx_die 2 "home directory not found: ${HOME_DIR}"
    id -u "${USERNAME}" >/dev/null 2>&1 || onx_die 2 "system user not found: ${USERNAME}"
fi

# ── Rollback wiring (only for fresh .ssh creation) ───────────────────────────
trap 'onx_rollback_run' ERR

SSH_DIR_CREATED=0
if [[ ! -d "${SSH_DIR}" ]]; then
    mkdir -p "${SSH_DIR}"
    SSH_DIR_CREATED=1
    onx_rollback_register "rm -rf '${SSH_DIR}' 2>/dev/null || true"
fi

# Enforce dir perms every time — sshd refuses 0775 etc.
chmod 0700 "${SSH_DIR}"
if [[ "${MOCK_MODE}" != "1" ]]; then
    chown "${USERNAME}:${USERNAME}" "${SSH_DIR}"
fi

# ── Append (idempotent) ──────────────────────────────────────────────────────
# Replace inline comment with the JSON-supplied comment if any was given.
LINE_TO_WRITE="${KEY}"
if [[ -n "${COMMENT}" ]]; then
    # Strip any existing trailing comment from the key (type + base64 + [comment])
    # and append our own. We keep type + body (first two fields).
    KEY_BODY="$(printf '%s' "${KEY}" | awk '{print $1" "$2}')"
    # Sanitise comment: no newlines, max 200 chars
    SAFE_COMMENT="$(printf '%s' "${COMMENT}" | tr -d '\r\n' | cut -c1-200)"
    LINE_TO_WRITE="${KEY_BODY} ${SAFE_COMMENT}"
fi

# Build the file if absent so we can grep without ENOENT noise.
if [[ ! -f "${AK_FILE}" ]]; then
    : > "${AK_FILE}"
    onx_rollback_register "rm -f '${AK_FILE}' 2>/dev/null || true"
fi
chmod 0600 "${AK_FILE}"

# Duplicate check by fingerprint — compute the fingerprint of every existing
# line and compare. Cheaper alternative (string equality on the body) misses
# the case where comment differs but the key bytes match.
ALREADY_PRESENT=0
if [[ -s "${AK_FILE}" ]]; then
    while IFS= read -r existing_line || [[ -n "${existing_line}" ]]; do
        # Skip blanks and comment lines
        case "${existing_line}" in
            ''|\#*) continue ;;
        esac
        # Compute fingerprint of this line
        EXIST_TMP="$(mktemp -t onx-ssh-exist.XXXXXX)"
        chmod 600 "${EXIST_TMP}"
        printf '%s\n' "${existing_line}" > "${EXIST_TMP}"
        EXIST_OUT="$(ssh-keygen -l -f "${EXIST_TMP}" 2>/dev/null || true)"
        rm -f "${EXIST_TMP}"
        EXIST_FP="$(printf '%s' "${EXIST_OUT}" | awk '{print $2}')"
        if [[ "${EXIST_FP}" == "${FINGERPRINT}" ]]; then
            ALREADY_PRESENT=1
            break
        fi
    done < "${AK_FILE}"
fi

if [[ "${ALREADY_PRESENT}" -eq 1 ]]; then
    # Disarm rollback — nothing to undo
    _ONX_ROLLBACK_STACK=()
    trap - ERR

    TOTAL_KEYS="$(grep -cE '^(ssh-|ecdsa-|sk-)' "${AK_FILE}" 2>/dev/null || echo 0)"
    onx_audit "onx-ssh" "duplicate-skip user=${USERNAME} fp=${FINGERPRINT}"
    onx_json_out \
        "username"    "${USERNAME}" \
        "fingerprint" "${FINGERPRINT}" \
        "key_type"    "${KEY_TYPE}" \
        "total_keys"  "${TOTAL_KEYS}" \
        "added"       "false"
    exit 0
fi

# Append + always finish with newline
if [[ -s "${AK_FILE}" ]] && [[ "$(tail -c1 "${AK_FILE}" 2>/dev/null || printf '')" != $'\n' ]]; then
    printf '\n' >> "${AK_FILE}"
fi
printf '%s\n' "${LINE_TO_WRITE}" >> "${AK_FILE}"

# Final perms / ownership pass (idempotent)
chmod 0600 "${AK_FILE}"
if [[ "${MOCK_MODE}" != "1" ]]; then
    chown "${USERNAME}:${USERNAME}" "${AK_FILE}"
fi

# Compute total live keys after append
TOTAL_KEYS="$(grep -cE '^(ssh-|ecdsa-|sk-)' "${AK_FILE}" 2>/dev/null || echo 0)"

onx_audit "onx-ssh" "add user=${USERNAME} fp=${FINGERPRINT} type=${KEY_TYPE} total=${TOTAL_KEYS}"

# Disarm rollback — success
_ONX_ROLLBACK_STACK=()
trap - ERR

onx_json_out \
    "username"    "${USERNAME}" \
    "fingerprint" "${FINGERPRINT}" \
    "key_type"    "${KEY_TYPE}" \
    "total_keys"  "${TOTAL_KEYS}" \
    "added"       "true"
