#!/usr/bin/env sh
set -eu

log() { echo "[vault-oidc] $*"; }

vault_cmd() {
  for attempt in 1 2 3 4 5 6; do
    set +e
    output="$(vault "$@" 2>&1)"
    status=$?
    set -e
    if [ "${status}" -eq 0 ]; then
      printf '%s' "${output}"
      return 0
    fi
    log "vault command failed; retrying (${attempt}/6)"
    sleep $((attempt * 2))
  done
  log "vault command failed; giving up"
  return 1
}

ensure_token() {
  if [ -n "${VAULT_TOKEN:-}" ]; then
    return
  fi
  role="${VAULT_K8S_ROLE:-vault}"
  jwt="$(cat /var/run/secrets/kubernetes.io/serviceaccount/token)"
  if ! VAULT_TOKEN="$(vault_cmd write -field=token auth/kubernetes/login role="${role}" jwt="${jwt}")"; then
    log "kubernetes auth login failed; set VAULT_TOKEN or fix role ${role}"
    exit 1
  fi
  export VAULT_TOKEN
}

if ! status_json="$(vault_cmd status -format=json)"; then
  log "vault status failed; check VAULT_ADDR and VAULT_TOKEN"
  exit 1
fi

if ! printf '%s' "${status_json}" | grep -q '"initialized":[[:space:]]*true'; then
  log "vault not initialized; skipping"
  exit 0
fi

if printf '%s' "${status_json}" | grep -q '"sealed":[[:space:]]*true'; then
  log "vault sealed; skipping"
  exit 0
fi

ensure_token

: "${VAULT_OIDC_DISCOVERY_URL:?set VAULT_OIDC_DISCOVERY_URL}"
: "${VAULT_OIDC_CLIENT_ID:?set VAULT_OIDC_CLIENT_ID}"
: "${VAULT_OIDC_CLIENT_SECRET:?set VAULT_OIDC_CLIENT_SECRET}"

default_role="${VAULT_OIDC_DEFAULT_ROLE:-admin}"
scopes="${VAULT_OIDC_SCOPES:-openid profile email groups}"
user_claim="${VAULT_OIDC_USER_CLAIM:-preferred_username}"
groups_claim="${VAULT_OIDC_GROUPS_CLAIM:-groups}"
redirect_uris="${VAULT_OIDC_REDIRECT_URIS:-https://secret.bstein.dev/ui/vault/auth/oidc/oidc/callback}"
bound_audiences="${VAULT_OIDC_BOUND_AUDIENCES:-${VAULT_OIDC_CLIENT_ID}}"
bound_claims_type="${VAULT_OIDC_BOUND_CLAIMS_TYPE:-string}"
bound_claims_type="$(printf '%s' "${bound_claims_type}" | tr -d '[:space:]')"
if [ -z "${bound_claims_type}" ] || [ "${bound_claims_type}" = "<novalue>" ]; then
  bound_claims_type="string"
fi

admin_group="${VAULT_OIDC_ADMIN_GROUP:-admin}"
admin_policies="${VAULT_OIDC_ADMIN_POLICIES:-default,vault-admin}"
dev_group="${VAULT_OIDC_DEV_GROUP:-dev}"
dev_policies="${VAULT_OIDC_DEV_POLICIES:-default,dev-kv}"
user_group="${VAULT_OIDC_USER_GROUP:-${dev_group}}"
user_policies="${VAULT_OIDC_USER_POLICIES:-${VAULT_OIDC_TOKEN_POLICIES:-${dev_policies}}}"

if ! vault_cmd auth list -format=json | grep -q '"oidc/"'; then
  log "enabling oidc auth method"
  vault_cmd auth enable oidc
fi

log "configuring oidc auth"
vault_cmd write auth/oidc/config \
  oidc_discovery_url="${VAULT_OIDC_DISCOVERY_URL}" \
  oidc_client_id="${VAULT_OIDC_CLIENT_ID}" \
  oidc_client_secret="${VAULT_OIDC_CLIENT_SECRET}" \
  default_role="${default_role}"

vault_cmd auth tune -listing-visibility=unauth oidc >/dev/null

build_bound_claims() {
  claim="$1"
  groups="$2"
  json="{\"${claim}\":["
  first=1
  old_ifs=$IFS
  IFS=,
  for item in $groups; do
    item="$(printf '%s' "$item" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"
    if [ -z "${item}" ]; then
      continue
    fi
    if [ "${first}" -eq 0 ]; then
      json="${json},"
    fi
    json="${json}\"${item}\""
    first=0
  done
  IFS=$old_ifs
  json="${json}]}"
  printf '%s' "${json}"
}

build_json_array() {
  items="$1"
  json="["
  first=1
  old_ifs=$IFS
  IFS=,
  for item in $items; do
    item="$(printf '%s' "$item" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"
    if [ -z "${item}" ]; then
      continue
    fi
    if [ "${first}" -eq 0 ]; then
      json="${json},"
    fi
    json="${json}\"${item}\""
    first=0
  done
  IFS=$old_ifs
  json="${json}]"
  printf '%s' "${json}"
}

configure_role() {
  role_name="$1"
  role_groups="$2"
  role_policies="$3"
  if [ -z "${role_name}" ] || [ -z "${role_groups}" ] || [ -z "${role_policies}" ]; then
    log "skipping role ${role_name} (missing groups or policies)"
    return
  fi
  claims="$(build_bound_claims "${groups_claim}" "${role_groups}")"
  scopes_csv="$(printf '%s' "${scopes}" | tr ' ' ',' | tr -s ',' | sed 's/^,//;s/,$//')"
  redirect_json="$(build_json_array "${redirect_uris}")"
  payload_file="$(mktemp)"
  cat > "${payload_file}" <<EOF
{
  "user_claim": "${user_claim}",
  "oidc_scopes": "${scopes_csv}",
  "token_policies": "${role_policies}",
  "bound_audiences": "${bound_audiences}",
  "bound_claims": ${claims},
  "bound_claims_type": "${bound_claims_type}",
  "groups_claim": "${groups_claim}",
  "allowed_redirect_uris": ${redirect_json}
}
EOF
  log "configuring oidc role ${role_name}"
  vault_cmd write "auth/oidc/role/${role_name}" @"${payload_file}"
  rm -f "${payload_file}"
}

configure_role "admin" "${admin_group}" "${admin_policies}"
configure_role "dev" "${dev_group}" "${dev_policies}"
configure_role "user" "${user_group}" "${user_policies}"
