#!/usr/bin/env bash
# openbao-session — SessionStart hook.
#
# Provisions a short-lived, policy-scoped OpenBao token for this interactive
# agent session. At session start it:
#   1. ensures the loopback SSH tunnel to the OpenBao private VPC listener is up,
#   2. caches the OpenBao TLS trust anchor locally,
#   3. mints a short-lived session token from a bootstrap credential,
#   4. makes it available to the session:
#        - Claude Code: writes `export OPENBAO_TOKEN=...` (and BAO_ADDR, BAO_CACERT,
#          BAO_TLS_SERVER_NAME) into $CLAUDE_ENV_FILE, which Claude Code sources
#          into every subsequent Bash command of the session.
#        - Other hosts (e.g. Codex), which cannot persist env from a hook: writes a
#          session-scoped env file (mode 600) and points the agent to it via
#          hookSpecificOutput.additionalContext — never the token value itself.
#
# Output: a single valid JSON object on stdout. On failure it still emits a
# benign `{"continue": true, "suppressOutput": true}` and exits 0 so the session
# starts cleanly; diagnostics go to stderr.
#
# One-time admin prerequisites and security notes: see README.md in this plugin.

set -euo pipefail

log() { echo "openbao-session: $*" >&2; }

# ---- configuration (env-overridable, defaults match the current production setup)
CONF_DIR="${OPENBAO_SESSION_CONF_DIR:-$HOME/.config/openbao}"
BIN_DIR="$CONF_DIR/bin"
LOCAL_PORT="${OPENBAO_SESSION_LOCAL_PORT:-8200}"
VPC_ADDR="${OPENBAO_SESSION_VPC_ADDR:-10.104.0.12}"
DROPLET="${OPENBAO_SESSION_DROPLET:-mbu-backup-jumphost}"
SSH_KEY="${OPENBAO_SESSION_SSH_KEY:-$HOME/.ssh/id_ed25519}"
HOST="${OPENBAO_SESSION_HOST:-}"                # optional explicit public IP
TLS_NAME="${OPENBAO_SESSION_TLS_NAME:-mbu-backup-jumphost}"
BAO_VER="${OPENBAO_SESSION_BAO_VER:-2.6.2}"
SESSION_POLICY="${OPENBAO_SESSION_POLICY:-claude-session}"
SESSION_TTL="${OPENBAO_SESSION_TTL:-4h}"
SESSION_ENTITY="${OPENBAO_SESSION_ENTITY:-}"   # OpenBao entity name; ties the session token to the user (enables humans/self via {{identity.entity.id}})
SECRETS_HELPER="${OPENBAO_SESSION_SECRETS_HELPER:-$HOME/.config/devops-secrets/manage.py}"
BOOTSTRAP_TOOL="${OPENBAO_SESSION_BOOTSTRAP_TOOL:-openbao}"
BOOTSTRAP_KEY="${OPENBAO_SESSION_BOOTSTRAP_KEY:-bootstrap_token}"
CACERT="$CONF_DIR/tls.crt"
BAO_BIN=""
BAO_ADDR_URL="https://127.0.0.1:${LOCAL_PORT}"
SESSION_TOKEN=""

benign_json() { printf '%s\n' '{"continue": true, "suppressOutput": true}'; }

ensure_bao() {
  if command -v bao >/dev/null 2>&1; then BAO_BIN="$(command -v bao)"; return 0; fi
  if [ -x "$BIN_DIR/bao" ]; then BAO_BIN="$BIN_DIR/bao"; return 0; fi
  mkdir -p "$BIN_DIR"
  local os arch asset url
  os="$(uname -s)"; arch="$(uname -m)"
  case "$os:$arch" in
    Darwin:arm64)               asset="openbao_${BAO_VER}_darwin_arm64.tar.gz" ;;
    Darwin:x86_64|Darwin:amd64) asset="openbao_${BAO_VER}_darwin_amd64.tar.gz" ;;
    Linux:arm64)                asset="openbao_${BAO_VER}_linux_arm64.tar.gz" ;;
    Linux:x86_64|Linux:amd64)   asset="openbao_${BAO_VER}_linux_amd64.tar.gz" ;;
    *) log "unsupported platform: $os/$arch"; return 1 ;;
  esac
  url="https://github.com/openbao/openbao/releases/download/v${BAO_VER}/${asset}"
  curl -fsSL -o "$CONF_DIR/bao.tar.gz" "$url" || { log "failed to download $asset"; return 1; }
  tar -xzf "$CONF_DIR/bao.tar.gz" -C "$BIN_DIR" bao || { log "failed to extract bao"; return 1; }
  chmod +x "$BIN_DIR/bao"
  rm -f "$CONF_DIR/bao.tar.gz"
  BAO_BIN="$BIN_DIR/bao"
}

resolve_host() {
  [ -n "$HOST" ] && return 0
  command -v doctl >/dev/null 2>&1 || { log "doctl missing and OPENBAO_SESSION_HOST unset"; return 1; }
  HOST="$(doctl compute droplet get "$DROPLET" --format PublicIPv4 --no-header 2>/dev/null | tr -d '[:space:]')"
  [[ "$HOST" =~ ^[0-9]{1,3}(\.[0-9]{1,3}){3}$ ]] || { log "cannot resolve $DROPLET public IPv4"; return 1; }
}

ensure_tunnel() {
  if nc -z 127.0.0.1 "$LOCAL_PORT" >/dev/null 2>&1; then return 0; fi
  [ -r "$SSH_KEY" ] || { log "SSH key not readable: $SSH_KEY"; return 1; }
  resolve_host || return 1
  ssh -f -N -i "$SSH_KEY" -o BatchMode=yes -o ExitOnForwardFailure=yes \
      -o ServerAliveInterval=30 -o ServerAliveCountMax=3 \
      -o StrictHostKeyChecking=accept-new \
      -L "127.0.0.1:${LOCAL_PORT}:${VPC_ADDR}:8200" "root@${HOST}" || return 1
  sleep 1
  nc -z 127.0.0.1 "$LOCAL_PORT" >/dev/null 2>&1 || { log "tunnel did not come up"; return 1; }
}

ensure_cacert() {
  [ -s "$CACERT" ] && return 0
  resolve_host || return 1
  mkdir -p "$CONF_DIR"
  scp -q -i "$SSH_KEY" -o BatchMode=yes -o StrictHostKeyChecking=accept-new \
      "root@${HOST}:/opt/openbao/tls/tls.crt" "$CACERT" || { log "failed to fetch CA cert"; return 1; }
}

read_bootstrap() {
  python3 "$SECRETS_HELPER" get "$BOOTSTRAP_TOOL" "$BOOTSTRAP_KEY" 2>/dev/null \
    || { log "bootstrap credential not found: $BOOTSTRAP_TOOL/$BOOTSTRAP_KEY (see README, one-time admin setup)"; return 1; }
}

mint_session_token() {
  local bootstrap json
  local -a entity_args=()
  bootstrap="$(read_bootstrap)" || return 1
  if [ -n "$SESSION_ENTITY" ]; then
    entity_args+=("-entity-alias=$SESSION_ENTITY")
  fi
  json="$(BAO_ADDR="$BAO_ADDR_URL" BAO_CACERT="$CACERT" BAO_TLS_SERVER_NAME="$TLS_NAME" BAO_TOKEN="$bootstrap" \
    "$BAO_BIN" token create -policy="$SESSION_POLICY" -ttl="$SESSION_TTL" -renewable \
    "${entity_args[@]}" \
    -display-name="openbao-session-$(hostname)" -format=json)" || { log "token create failed"; return 1; }
  SESSION_TOKEN="$(printf '%s' "$json" | jq -r '.auth.client_token // empty')"
  [ -n "$SESSION_TOKEN" ] || { log "no client_token in token create response"; return 1; }
}

session_id_from_stdin() {
  local input
  input="$(cat 2>/dev/null || true)"
  [ -n "$input" ] && command -v jq >/dev/null 2>&1 \
    && printf '%s' "$input" | jq -r '.session_id // empty' 2>/dev/null
}

write_env_exports() {
  printf 'export OPENBAO_TOKEN=%q\n' "$SESSION_TOKEN"
  printf 'export BAO_ADDR=%q\n' "$BAO_ADDR_URL"
  printf 'export BAO_CACERT=%q\n' "$CACERT"
  printf 'export BAO_TLS_SERVER_NAME=%q\n' "$TLS_NAME"
}

inject_env() {
  if [ -n "${CLAUDE_ENV_FILE:-}" ]; then
    # Claude Code: persist exports for every subsequent Bash command this session.
    write_env_exports >> "$CLAUDE_ENV_FILE" || { log "failed to write CLAUDE_ENV_FILE"; return 1; }
    benign_json
    return 0
  fi
  # Other hosts (Codex etc.): hooks cannot persist env; write a session-scoped
  # env file and point the agent to it. Never put the token value in context.
  local sid envfile
  sid="$(session_id_from_stdin)"
  envfile="${TMPDIR:-/tmp}/openbao-session.env.${sid:-$$}"
  write_env_exports > "$envfile" || return 1
  chmod 600 "$envfile"
  printf '{"continue":true,"suppressOutput":true,"hookSpecificOutput":{"additionalContext":"OpenBao session token ready. In a Bash tool, load it without printing the value: source %s"}}\n' "$envfile"
}

run() {
  ensure_bao || return 1
  ensure_tunnel || return 1
  ensure_cacert || return 1
  mint_session_token || return 1
  inject_env || return 1
}

main() {
  local out rc=0
  out="$(run)" || rc=1
  if [ "$rc" -eq 0 ]; then
    printf '%s\n' "$out"
  else
    log "provisioning failed; session continues without an OpenBao session token"
    benign_json
  fi
  exit 0
}

main "$@"