#!/usr/bin/env bash
# CASAN H5 — HashiCorp Vault KMS helper
#
# Provides Vault-backed signing operations as a drop-in replacement for
# direct openssl key-file usage.  Scripts detect KMS availability via:
#   [[ -n "${VAULT_ADDR:-}" && -n "${VAULT_TOKEN:-}" ]]
#
# Usage (direct):
#   vault-kms.sh sign   <file> <output.sig> [key_name]
#   vault-kms.sh verify <file> <sig_file>   [key_name]
#   vault-kms.sh pubkey <output.pem>        [key_name]
#   vault-kms.sh status
#
# Usage (sourced):
#   source vault-kms.sh
#   vault_kms_sign   <file> <output.sig> [key_name]
#   vault_kms_pubkey <output.pem>        [key_name]
#   vault_kms_status
#
# Environment:
#   VAULT_ADDR   — e.g. http://vault:8200 or http://161.33.139.73:8200
#   VAULT_TOKEN  — root token or policy token with transit/sign/* capability

set -euo pipefail

: "${VAULT_ADDR:?VAULT_ADDR must be set}"
: "${VAULT_TOKEN:?VAULT_TOKEN must be set}"

_VAULT_DEFAULT_KEY="casan-policy-key"
_VAULT_AUDIT_KEY="casan-audit-key"

# ── Ensure required tools ──────────────────────────────────────────────────
_check_deps() {
  for cmd in curl python3; do
    command -v "$cmd" >/dev/null 2>&1 || { echo "vault-kms: required: $cmd" >&2; exit 1; }
  done
}

# ── Transit: ensure key exists ────────────────────────────────────────────
vault_kms_ensure_key() {
  local key="${1:-$_VAULT_DEFAULT_KEY}"
  local status
  status=$(curl -sf \
    -H "X-Vault-Token: $VAULT_TOKEN" \
    "$VAULT_ADDR/v1/transit/keys/$key" 2>/dev/null | \
    python3 -c "import sys,json; d=json.load(sys.stdin); print('ok' if 'data' in d else 'missing')" 2>/dev/null || echo "missing")

  if [[ "$status" != "ok" ]]; then
    curl -sf -X POST \
      -H "X-Vault-Token: $VAULT_TOKEN" \
      -H "Content-Type: application/json" \
      -d '{"type":"rsa-2048","exportable":false,"allow_plaintext_backup":false}' \
      "$VAULT_ADDR/v1/transit/keys/$key" >/dev/null
    echo "vault-kms: created transit key: $key" >&2
  fi
}

# ── Sign a file ───────────────────────────────────────────────────────────
# Produces a DER binary signature compatible with:
#   openssl dgst -sha256 -verify pub.pem -signature <output.sig> <file>
vault_kms_sign() {
  local file="$1" output="$2" key="${3:-$_VAULT_DEFAULT_KEY}"
  _check_deps
  vault_kms_ensure_key "$key"

  # Base64-encode the file content for the Vault API payload
  local input_b64
  input_b64=$(base64 -w0 < "$file" 2>/dev/null || base64 < "$file")

  # Call Vault Transit sign — PKCS#1 v1.5 + SHA-256 (compatible with openssl verify)
  local response
  response=$(curl -sf -X POST \
    -H "X-Vault-Token: $VAULT_TOKEN" \
    -H "Content-Type: application/json" \
    -d "{\"input\":\"$input_b64\",\"hash_algorithm\":\"sha2-256\",\"signature_algorithm\":\"pkcs1v15\",\"prehashed\":false}" \
    "$VAULT_ADDR/v1/transit/sign/$key") || {
      echo "vault-kms: sign request failed (key=$key)" >&2; exit 1
    }

  # Extract signature and strip the "vault:v1:" prefix → raw base64 DER bytes
  local sig_b64
  sig_b64=$(printf '%s' "$response" | python3 -c "
import sys, json
d = json.load(sys.stdin)
sig = d['data']['signature']
# vault:v1:<base64>  →  keep only base64 part
print(sig.split(':')[-1])
")

  # Decode base64 → binary DER file (identical format to openssl -sign output)
  printf '%s' "$sig_b64" | base64 -d > "$output"
  echo "vault-kms: signed $file → $output (key=$key anchor=vault-kms)" >&2
}

# ── Export public key ─────────────────────────────────────────────────────
# Writes the RSA public key as PEM so openssl dgst -verify still works.
vault_kms_pubkey() {
  local output="$1" key="${2:-$_VAULT_DEFAULT_KEY}"
  _check_deps
  vault_kms_ensure_key "$key"

  local response
  response=$(curl -sf \
    -H "X-Vault-Token: $VAULT_TOKEN" \
    "$VAULT_ADDR/v1/transit/keys/$key") || {
      echo "vault-kms: pubkey fetch failed (key=$key)" >&2; exit 1
    }

  local tmp_resp
  tmp_resp=$(mktemp)
  printf '%s' "$response" > "$tmp_resp"
  python3 - "$output" "$tmp_resp" <<'PY'
import sys, json
output = sys.argv[1]
with open(sys.argv[2]) as f:
    d = json.loads(f.read())
keys = d["data"]["keys"]
# keys is a dict; pick the latest version
latest = max(keys.keys(), key=lambda k: int(k))
pub = keys[latest]["public_key"]
with open(output, "w") as f:
    f.write(pub if pub.endswith("\n") else pub + "\n")
print(f"vault-kms: public key written to {output}", file=sys.stderr)
PY
  rm -f "$tmp_resp"
}

# ── Verify a signature ────────────────────────────────────────────────────
# Uses openssl with the public key exported from Vault.
vault_kms_verify() {
  local file="$1" sig_file="$2" key="${3:-$_VAULT_DEFAULT_KEY}"
  local tmp_pub
  tmp_pub=$(mktemp /tmp/vault-pub-XXXXX.pem)
  trap 'rm -f "$tmp_pub"' RETURN
  vault_kms_pubkey "$tmp_pub" "$key"
  openssl dgst -sha256 -verify "$tmp_pub" -signature "$sig_file" "$file" >/dev/null
}

# ── Connectivity check ────────────────────────────────────────────────────
vault_kms_status() {
  if curl -sf "$VAULT_ADDR/v1/sys/health" >/dev/null 2>&1; then
    echo "VAULT_KMS_AVAILABLE addr=$VAULT_ADDR"
    return 0
  else
    echo "VAULT_KMS_UNAVAILABLE addr=${VAULT_ADDR:-unset}"
    return 1
  fi
}

# ── Enable transit engine (idempotent) ────────────────────────────────────
vault_kms_enable_transit() {
  curl -sf -X POST \
    -H "X-Vault-Token: $VAULT_TOKEN" \
    -H "Content-Type: application/json" \
    -d '{"type":"transit"}' \
    "$VAULT_ADDR/v1/sys/mounts/transit" >/dev/null 2>&1 || true
}

# ── CLI entrypoint ────────────────────────────────────────────────────────
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  CMD="${1:-status}"
  shift || true
  case "$CMD" in
    sign)    vault_kms_sign   "$@" ;;
    pubkey)  vault_kms_pubkey "$@" ;;
    verify)  vault_kms_verify "$@" ;;
    status)  vault_kms_status ;;
    enable-transit) vault_kms_enable_transit ;;
    ensure-key)     vault_kms_ensure_key "${1:-}" ;;
    *)
      echo "Usage: vault-kms.sh {sign|pubkey|verify|status|enable-transit|ensure-key}" >&2
      exit 64
      ;;
  esac
fi