#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
VAULT_ADDR="${VAULT_ADDR:-http://127.0.0.1:8200}"
VAULT_INIT_OUTPUT="${VAULT_INIT_OUTPUT:-${REPO_ROOT}/.local/vault/dev-k3s-init.json}"

for cmd in jq rg vault; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "$cmd is required" >&2
    exit 1
  fi
done

export VAULT_ADDR

status_json() {
  local output=""
  local status_code=0

  set +e
  output="$(vault status -format=json 2>/dev/null)"
  status_code=$?
  set -e

  if [[ "$status_code" -ne 0 && "$status_code" -ne 2 ]]; then
    echo "Vault is not reachable at ${VAULT_ADDR}." >&2
    return "$status_code"
  fi
  printf '%s\n' "$output"
}

unseal() {
  local status=""
  local unseal_key=""

  status="$(status_json)"
  if [[ "$(jq -r '.sealed' <<<"$status")" == "false" ]]; then
    echo "Vault is already unsealed."
    return
  fi
  if [[ ! -f "$VAULT_INIT_OUTPUT" ]]; then
    echo "Init material is unavailable: ${VAULT_INIT_OUTPUT}" >&2
    exit 1
  fi

  unseal_key="$(jq -er '.unseal_keys_b64[0]' "$VAULT_INIT_OUTPUT")"
  vault operator unseal "$unseal_key" >/dev/null
  echo "Vault is unsealed."
}

init() {
  local status=""

  status="$(status_json)"
  if [[ "$(jq -r '.initialized' <<<"$status")" == "true" ]]; then
    echo "Refusing initialization: Vault is already initialized." >&2
    exit 1
  fi
  if [[ -e "$VAULT_INIT_OUTPUT" ]]; then
    echo "Refusing to overwrite existing init material: ${VAULT_INIT_OUTPUT}" >&2
    exit 1
  fi

  umask 077
  mkdir -p "$(dirname "$VAULT_INIT_OUTPUT")"
  vault operator init \
    -key-shares=1 \
    -key-threshold=1 \
    -format=json >"$VAULT_INIT_OUTPUT"
  chmod 0600 "$VAULT_INIT_OUTPUT"
  unseal

  echo "Dev Vault was initialized with a dev-only 1-of-1 Shamir key."
  echo "Move ${VAULT_INIT_OUTPUT} to encrypted custody before continuing."
}

revoke_root() {
  local database_lookup=""
  local database_token="${VAULT_DATABASE_TOKEN:-}"
  local root_token=""
  local temporary=""
  local workloads_lookup=""
  local workloads_token="${VAULT_WORKLOADS_TOKEN:-}"

  if [[ ! -f "$VAULT_INIT_OUTPUT" ]]; then
    echo "Init material is unavailable: ${VAULT_INIT_OUTPUT}" >&2
    exit 1
  fi
  if [[ -z "$workloads_token" || -z "$database_token" ]]; then
    echo "Refusing root revocation: VAULT_WORKLOADS_TOKEN and VAULT_DATABASE_TOKEN are required." >&2
    exit 1
  fi

  if ! VAULT_TOKEN="$workloads_token" \
    vault token capabilities sys/policies/acl/auth-server-dev |
    rg -q '(^|,|[[:space:]])update($|,|[[:space:]])'; then
    echo "Refusing root revocation: the workloads replacement token failed its capability check." >&2
    exit 1
  fi
  if ! VAULT_TOKEN="$database_token" \
    vault token capabilities database/config/auth-system-postgres-dev |
    rg -q '(^|,|[[:space:]])update($|,|[[:space:]])'; then
    echo "Refusing root revocation: the database replacement token failed its capability check." >&2
    exit 1
  fi
  if ! workloads_lookup="$(
    VAULT_TOKEN="$workloads_token" vault token lookup -format=json
  )"; then
    echo "Refusing root revocation: the workloads replacement token cannot look itself up." >&2
    exit 1
  fi
  if ! database_lookup="$(
    VAULT_TOKEN="$database_token" vault token lookup -format=json
  )"; then
    echo "Refusing root revocation: the database replacement token cannot look itself up." >&2
    exit 1
  fi
  if ! jq -e '.data.policies | type == "array"' <<<"$workloads_lookup" >/dev/null ||
    ! jq -e '.data.policies | type == "array"' <<<"$database_lookup" >/dev/null; then
    echo "Refusing root revocation: a replacement token returned an invalid lookup response." >&2
    exit 1
  fi
  if jq -e '.data.policies | index("root") != null' <<<"$workloads_lookup" >/dev/null ||
    jq -e '.data.policies | index("root") != null' <<<"$database_lookup" >/dev/null; then
    echo "Refusing root revocation: a replacement token unexpectedly carries the root policy." >&2
    exit 1
  fi

  root_token="$(jq -er '.root_token' "$VAULT_INIT_OUTPUT")"
  VAULT_TOKEN="$root_token" vault token revoke -self

  temporary="$(mktemp "${VAULT_INIT_OUTPUT}.XXXXXX")"
  jq 'del(.root_token)' "$VAULT_INIT_OUTPUT" >"$temporary"
  chmod 0600 "$temporary"
  mv "$temporary" "$VAULT_INIT_OUTPUT"
  echo "The initial root token was revoked and removed from the local init file."
}

case "${1:-}" in
  init)
    init
    ;;
  unseal)
    unseal
    ;;
  revoke-root)
    revoke_root
    ;;
  *)
    echo "Usage: $0 <init|unseal|revoke-root>" >&2
    exit 1
    ;;
esac