#!/usr/bin/env bash
# =============================================================================
# onx-ssh-key-remove — Drop an SSH public key from a user's authorized_keys.
#
# Input (stdin JSON):
#   {
#     "username":    "onx_acme01",                -- required
#     "fingerprint": "SHA256:abc..."              -- one-of: fingerprint OR key
#     "key":         "ssh-rsa AAAA... user@host"  -- one-of: fingerprint OR key
#   }
#
# Output (stdout JSON):
#   {
#     "username":      "onx_acme01",
#     "removed":       true,
#     "remaining_keys": 2
#   }
#
# Exit codes: 0=ok 1=invalid 2=preflight 3=exec 4=rolled-back 5=rollback-fail
# =============================================================================

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")
FINGERPRINT=$(onx_json_field "fingerprint" "")
KEY=$(onx_json_field         "key" "")

onx_validate_username "${USERNAME}"

if [[ -z "${FINGERPRINT}" && -z "${KEY}" ]]; then
    onx_die 1 "either 'fingerprint' or 'key' is required"
fi

# If only `key` was supplied, derive the fingerprint
if [[ -z "${FINGERPRINT}" ]]; then
    KEY="$(printf '%s' "${KEY}" | tr -d '\r' | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')"
    TMP_KEY="$(mktemp -t onx-ssh-rm.XXXXXX)"
    chmod 600 "${TMP_KEY}"
    trap '[[ -n "${TMP_KEY:-}" ]] && rm -f "${TMP_KEY}" 2>/dev/null || true' EXIT
    printf '%s\n' "${KEY}" > "${TMP_KEY}"
    FINGERPRINT="$(ssh-keygen -l -f "${TMP_KEY}" 2>/dev/null | awk '{print $2}')"
    [[ "${FINGERPRINT}" =~ ^SHA256: ]] || onx_die 1 "could not derive fingerprint from supplied key"
fi

case "${FINGERPRINT}" in
    SHA256:*) : ;;
    *) onx_die 1 "fingerprint must be SHA256:... format" ;;
esac

HOME_DIR="/home/${USERNAME}"
AK_FILE="${HOME_DIR}/.ssh/authorized_keys"

if [[ ! -f "${AK_FILE}" ]]; then
    # Nothing to remove — return idempotent success with 0 remaining
    onx_audit "onx-ssh" "remove-noop user=${USERNAME} fp=${FINGERPRINT} (no authorized_keys file)"
    onx_json_out \
        "username"       "${USERNAME}" \
        "removed"        "false" \
        "remaining_keys" "0"
    exit 0
fi

# Backup so we can roll back
BACKUP_FILE="${AK_FILE}.onx-bak.$$"
cp -p "${AK_FILE}" "${BACKUP_FILE}"
trap 'cp -p "${BACKUP_FILE}" "${AK_FILE}" 2>/dev/null; rm -f "${BACKUP_FILE}" 2>/dev/null || true' ERR

# Walk every line, recompute fingerprint, drop matches
NEW_FILE="$(mktemp -t onx-ak-new.XXXXXX)"
chmod 600 "${NEW_FILE}"
REMOVED_COUNT=0

while IFS= read -r line || [[ -n "${line}" ]]; do
    # Pass blanks and pure comments through untouched
    case "${line}" in
        ''|\#*)
            printf '%s\n' "${line}" >> "${NEW_FILE}"
            continue
            ;;
    esac
    # Compute this line's fingerprint
    TMP_FP="$(mktemp -t onx-fp.XXXXXX)"
    chmod 600 "${TMP_FP}"
    printf '%s\n' "${line}" > "${TMP_FP}"
    LINE_FP="$(ssh-keygen -l -f "${TMP_FP}" 2>/dev/null | awk '{print $2}' || true)"
    rm -f "${TMP_FP}"

    if [[ "${LINE_FP}" == "${FINGERPRINT}" ]]; then
        REMOVED_COUNT=$((REMOVED_COUNT + 1))
        continue
    fi
    printf '%s\n' "${line}" >> "${NEW_FILE}"
done < "${AK_FILE}"

# Move into place + enforce perms
mv "${NEW_FILE}" "${AK_FILE}"
chmod 0600 "${AK_FILE}"
if [[ "${MOCK_MODE}" != "1" ]]; then
    chown "${USERNAME}:${USERNAME}" "${AK_FILE}" 2>/dev/null || true
fi

# Drop the rollback backup — we succeeded
rm -f "${BACKUP_FILE}" 2>/dev/null || true
trap - ERR

REMAINING_KEYS="$(grep -cE '^(ssh-|ecdsa-|sk-)' "${AK_FILE}" 2>/dev/null || echo 0)"

if [[ "${REMOVED_COUNT}" -eq 0 ]]; then
    onx_audit "onx-ssh" "remove-miss user=${USERNAME} fp=${FINGERPRINT}"
    onx_json_out \
        "username"       "${USERNAME}" \
        "removed"        "false" \
        "remaining_keys" "${REMAINING_KEYS}"
    exit 0
fi

onx_audit "onx-ssh" "remove user=${USERNAME} fp=${FINGERPRINT} count=${REMOVED_COUNT} remaining=${REMAINING_KEYS}"
onx_json_out \
    "username"       "${USERNAME}" \
    "removed"        "true" \
    "remaining_keys" "${REMAINING_KEYS}"
