#!/usr/bin/env bash
# vault/startup.sh
#
# Boot script for Vault HA cluster nodes on GCE.
# Idempotent — safe to re-run on VM restart.
#
# Required env vars (set by the GCE metadata startup-script):
#   VAULT_NODE_ID  — e.g. "vault-node-1"
#   VAULT_VERSION  — e.g. "1.19.2"
#
# After first boot:
#   1. On vault-node-1: vault operator init
#      → save root token and recovery keys in a secure location (Vault bootstrap
#        requires them; the recovery keys are for KMS-sealed Vault)
#   2. Nodes 2 and 3 auto-join via retry_join once initialized
#
# Ops access:
#   gcloud compute ssh vault-node-1 --zone=us-central1-a --tunnel-through-iap

set -euxo pipefail
exec > >(tee /var/log/vault-bootstrap.log) 2>&1

: "${VAULT_NODE_ID:?VAULT_NODE_ID must be set}"
: "${VAULT_VERSION:?VAULT_VERSION must be set}"

# ── Install Vault ─────────────────────────────────────────────────────────────

apt-get update -y
apt-get install -y curl ca-certificates gnupg lsb-release unzip

# HashiCorp GPG key + repo
curl -fsSL https://apt.releases.hashicorp.com/gpg \
  | gpg --dearmor -o /usr/share/keyrings/hashicorp-archive-keyring.gpg

echo "deb [signed-by=/usr/share/keyrings/hashicorp-archive-keyring.gpg] \
  https://apt.releases.hashicorp.com $(lsb_release -cs) main" \
  > /etc/apt/sources.list.d/hashicorp.list

apt-get update -y
apt-get install -y vault="${VAULT_VERSION}-1"

# Give Vault cap to lock memory (prevents secrets being swapped to disk)
setcap cap_ipc_lock=+ep /usr/bin/vault

# ── Mount and prepare the data disk ──────────────────────────────────────────
# The data disk is attached as /dev/disk/by-id/google-vault-data
# First boot: format it. Subsequent boots: just mount it.

DATA_DEVICE="/dev/disk/by-id/google-vault-data"
DATA_MOUNT="/opt/vault/data"
mkdir -p "${DATA_MOUNT}"

# Only format if not already formatted (idempotent)
if ! blkid "${DATA_DEVICE}" | grep -q ext4; then
  mkfs.ext4 -F "${DATA_DEVICE}"
fi

# Mount (skip if already mounted)
if ! mountpoint -q "${DATA_MOUNT}"; then
  mount "${DATA_DEVICE}" "${DATA_MOUNT}"
fi

# Persist in fstab
DEVICE_UUID="$(blkid -s UUID -o value "${DATA_DEVICE}")"
if ! grep -q "${DEVICE_UUID}" /etc/fstab; then
  echo "UUID=${DEVICE_UUID} ${DATA_MOUNT} ext4 defaults,noatime 0 2" >> /etc/fstab
fi

# ── Derive internal IPs for retry_join ───────────────────────────────────────
# GCE metadata server gives us the IPs of the other nodes by instance name.
# We query the Compute API via the attached SA (ADC — no key file needed).
# For the initial single-node bootstrap, retry_join failures are harmless
# (Raft simply starts as a single-node cluster until peers join).

fetch_internal_ip() {
  local instance_name="$1"
  local zone="$2"
  gcloud compute instances describe "${instance_name}" \
    --zone="${zone}" \
    --format="get(networkInterfaces[0].networkIP)" \
    2>/dev/null || echo ""
}

# Map of all Vault nodes: name → zone
declare -A NODE_ZONES
NODE_ZONES["vault-node-1"]="us-central1-a"
NODE_ZONES["vault-node-2"]="us-central1-b"
NODE_ZONES["vault-node-3"]="us-central1-c"

# Build retry_join stanzas for all nodes other than self
RETRY_JOIN_STANZAS=""
for node in "vault-node-1" "vault-node-2" "vault-node-3"; do
  if [[ "${node}" == "${VAULT_NODE_ID}" ]]; then
    continue
  fi
  zone="${NODE_ZONES[$node]}"
  ip="$(fetch_internal_ip "${node}" "${zone}")"
  if [[ -n "${ip}" ]]; then
    RETRY_JOIN_STANZAS="${RETRY_JOIN_STANZAS}
  retry_join {
    leader_api_addr = \"http://${ip}:8200\"
  }"
  fi
done

# ── Get the internal IP of this node ─────────────────────────────────────────
MY_INTERNAL_IP="$(curl -sf -H "Metadata-Flavor: Google" \
  "http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/ip")"

# ── Write Vault config ────────────────────────────────────────────────────────

mkdir -p /etc/vault.d
cat > /etc/vault.d/vault.hcl <<EOF
ui = true

listener "tcp" {
  address       = "0.0.0.0:8200"
  cluster_address = "0.0.0.0:8201"
  # TLS terminated at GCP LB — plain HTTP internally is fine
  tls_disable   = true
}

storage "raft" {
  path    = "/opt/vault/data"
  node_id = "${VAULT_NODE_ID}"
${RETRY_JOIN_STANZAS}
}

seal "gcpckms" {
  project    = "neuron-785695"
  region     = "global"
  key_ring   = "vault"
  crypto_key = "vault-unseal"
}

api_addr     = "http://${MY_INTERNAL_IP}:8200"
cluster_addr = "http://${MY_INTERNAL_IP}:8201"

# Telemetry for Cloud Monitoring
telemetry {
  prometheus_retention_time = "30s"
  disable_hostname          = false
}
EOF

chmod 640 /etc/vault.d/vault.hcl

# ── Create vault user and fix permissions ─────────────────────────────────────

id vault &>/dev/null || useradd --system --home /etc/vault.d --shell /bin/false vault
chown -R vault:vault /etc/vault.d /opt/vault/data

# ── Systemd unit ──────────────────────────────────────────────────────────────

cat > /etc/systemd/system/vault.service <<'EOF'
[Unit]
Description=HashiCorp Vault — a tool for managing secrets
Documentation=https://developer.hashicorp.com/vault/docs
Requires=network-online.target
After=network-online.target
ConditionFileNotEmpty=/etc/vault.d/vault.hcl
StartLimitIntervalSec=60
StartLimitBurst=3

[Service]
Type=notify
User=vault
Group=vault
ProtectSystem=full
ProtectHome=read-only
PrivateTmp=yes
PrivateDevices=yes
SecureBits=keep-caps
AmbientCapabilities=CAP_IPC_LOCK
CapabilityBoundingSet=CAP_SYSLOG CAP_IPC_LOCK
NoNewPrivileges=yes
ExecStart=/usr/bin/vault server -config=/etc/vault.d/vault.hcl
ExecReload=/bin/kill --signal HUP $MAINPID
KillMode=process
KillSignal=SIGINT
Restart=on-failure
RestartSec=5
TimeoutStopSec=30
LimitNOFILE=65536
LimitMEMLOCK=infinity

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable vault

# Start or restart (restart is idempotent — re-reads config on VM metadata updates)
if systemctl is-active vault &>/dev/null; then
  systemctl restart vault
else
  systemctl start vault
fi

echo "Vault ${VAULT_NODE_ID} started successfully"
echo "Run 'vault operator init' on vault-node-1 to bootstrap the cluster"