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

PROLE_USER="${PROLE_USER:-prole}"
TOKEN="${OPENBAO_TOKEN:-}"
TOKEN_FILE="${OPENBAO_TOKEN_FILE:-/run/secrets/openbao_token}"

if [[ -z "$TOKEN" && -f "$TOKEN_FILE" ]]; then
  TOKEN=$(cat "$TOKEN_FILE")
fi

if [[ -z "$TOKEN" ]]; then
  exit 0
fi

NAMESPACE="${NAMESPACE:-}"
if [[ -z "$NAMESPACE" && -f /var/run/secrets/kubernetes.io/serviceaccount/namespace ]]; then
  NAMESPACE=$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace)
fi

OPENBAO_ADDR="${OPENBAO_ADDR:-${PROLE_OPENBAO_URL:-}}"
if [[ -z "$OPENBAO_ADDR" && -n "$NAMESPACE" ]]; then
  OPENBAO_ADDR="http://openbao.${NAMESPACE}.svc.cluster.local:8200"
fi

if [[ -z "$OPENBAO_ADDR" ]]; then
  exit 0
fi

OPENBAO_SSH_MOUNT="${OPENBAO_SSH_MOUNT:-kv}"
OPENBAO_SSH_PATH="${OPENBAO_SSH_PATH:-prole/${NAMESPACE}/admin}"

if [[ -z "$OPENBAO_SSH_PATH" || "$OPENBAO_SSH_PATH" == "prole//admin" ]]; then
  exit 0
fi

response=$(curl -fsS -H "X-Vault-Token: $TOKEN" \
  "$OPENBAO_ADDR/v1/${OPENBAO_SSH_MOUNT}/data/${OPENBAO_SSH_PATH}" 2>/dev/null || true)

if [[ -z "$response" ]]; then
  exit 0
fi

ssh_key=$(python3 - <<'PY' <<<"$response"
import base64
import json
import sys

try:
    payload = json.load(sys.stdin)
except Exception:
    sys.exit(0)

data = payload.get("data", {}).get("data", {})
key_b64 = data.get("admin_public_key_b64")
if key_b64:
    try:
        print(base64.b64decode(key_b64).decode().strip())
        sys.exit(0)
    except Exception:
        pass

for k in ("admin_public_key", "ssh_public_key", "public_key"):
    val = data.get(k)
    if val:
        print(str(val).strip())
        break
PY
)

if [[ -z "$ssh_key" ]]; then
  exit 0
fi

if ! id -u "$PROLE_USER" >/dev/null 2>&1; then
  exit 0
fi

ssh_dir="/home/$PROLE_USER/.ssh"
auth_keys="${ssh_dir}/authorized_keys"

mkdir -p "$ssh_dir"
touch "$auth_keys"

if ! grep -qxF "$ssh_key" "$auth_keys"; then
  printf '%s\n' "$ssh_key" >> "$auth_keys"
fi

chown -R "$PROLE_USER:$PROLE_USER" "$ssh_dir"
chmod 700 "$ssh_dir"
chmod 600 "$auth_keys"