#!/bin/bash
# qxvault-2fa-enable — Enable 2FA (second-pass unlock) for a SecureVM-encrypted VM
# Part of proxmox-kms-bridge (45Drives SecureVM for Proxmox)
#
# Usage: qxvault-2fa-enable <vmid> <mode> [passphrase] [totp_secret_b32]
#   mode: passphrase | totp
#
# In passphrase mode:
#   - Derives a combined key from vault_key + user_passphrase using PBKDF2/HKDF
#   - Adds the combined key as a new LUKS keyslot
#   - Removes the vault-only keyslot
#   - The passphrase is required alongside the vault key to unlock
#
# In TOTP mode:
#   - Generates a TOTP secret, stores it encrypted in /etc/qxvault/2fa/<vmid>.totp
#   - The hookscript blocks auto-unlock; admin must provide TOTP code via UI
#   - The vault key alone is cryptographically sufficient (TOTP is an authorization gate)
#
# Security: passphrase is read from stdin or passed as arg (stdin preferred).
set -euo pipefail

CONF_FILE="/etc/qxvault/qxvault.conf"
TFA_DIR="/etc/qxvault/2fa"
VAULT_DMKEY="/usr/libexec/vault-dmkey"

# ── Defaults ───────────────────────────────────────────────────────
LUKS_IMAGE_DIR="/var/lib/vz/images"
LUKS_IMAGE_TEMPLATE="vm-{vmid}-disk-qxvault.luks"
LUKS_MAPPER_TEMPLATE="vm{vmid}_crypt"

if [ -f "$CONF_FILE" ]; then
  # shellcheck source=/dev/null
  . "$CONF_FILE"
fi

# ── Guards ─────────────────────────────────────────────────────────
[ "$(id -u)" -eq 0 ] || { echo "error: must run as root" >&2; exit 1; }

VMID="${1:-}"
MODE="${2:-}"
PASSPHRASE="${3:-}"
TOTP_SECRET_OVERRIDE="${4:-}"

[ -n "$VMID" ] || { echo "usage: qxvault-2fa-enable <vmid> <mode> [passphrase]" >&2; exit 1; }
[[ "$VMID" =~ ^[0-9]+$ ]] || { echo "error: VMID must be numeric" >&2; exit 1; }
[[ "$MODE" =~ ^(passphrase|totp)$ ]] || { echo "error: mode must be 'passphrase' or 'totp'" >&2; exit 1; }

IMAGE="${LUKS_IMAGE_DIR}/${VMID}/${LUKS_IMAGE_TEMPLATE//\{vmid\}/$VMID}"
MAPPER="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$VMID}"

[ -f "$IMAGE" ] || { echo "error: LUKS image not found: $IMAGE" >&2; exit 1; }

# Check VM is stopped
if command -v qm >/dev/null 2>&1; then
  vm_status=$(qm status "$VMID" 2>/dev/null || echo "stopped")
  if echo "$vm_status" | grep -q running; then
    echo "error: VM $VMID is running. Stop it first." >&2
    exit 1
  fi
fi

# Check mapper is closed
if cryptsetup status "$MAPPER" >/dev/null 2>&1; then
  echo "error: LUKS mapper $MAPPER is open. Close it first." >&2
  exit 1
fi

# Check not already 2FA-enabled
if [ -f "${TFA_DIR}/${VMID}.conf" ]; then
  echo "error: 2FA is already enabled for VMID $VMID. Disable first." >&2
  exit 1
fi

# ── Create 2FA directory ───────────────────────────────────────────
mkdir -p "$TFA_DIR"
chmod 0700 "$TFA_DIR"

# ── Retrieve current vault key ─────────────────────────────────────
echo "Retrieving vault key for VMID $VMID..."
KEY_B64=$("$VAULT_DMKEY" get-key --name "$VMID" 2>/dev/null)
[ -n "$KEY_B64" ] || { echo "error: failed to retrieve vault key" >&2; exit 1; }

# Write vault key to temp file (binary data can't be stored in shell variables)
VAULT_KEYFILE=$(mktemp /dev/shm/qxvault-vkey-XXXXXX)
chmod 0600 "$VAULT_KEYFILE"
echo "$KEY_B64" | base64 -d > "$VAULT_KEYFILE"
trap 'rm -f "$VAULT_KEYFILE"' EXIT

# Verify the vault key opens the LUKS image
echo "Verifying vault key against LUKS header..."
if ! cryptsetup open --test-passphrase --type luks --key-file "$VAULT_KEYFILE" "$IMAGE"; then
  echo "error: vault key does not match LUKS header" >&2
  rm -f "$VAULT_KEYFILE"
  exit 1
fi

# ── Mode-specific setup ────────────────────────────────────────────
case "$MODE" in
  passphrase)
    # Read passphrase from stdin if not passed as argument
    if [ -z "$PASSPHRASE" ]; then
      echo "Enter 2FA passphrase for VM $VMID:" >&2
      read -rs PASSPHRASE
      echo >&2
      [ -n "$PASSPHRASE" ] || { echo "error: passphrase cannot be empty" >&2; exit 1; }
      echo "Confirm passphrase:" >&2
      read -rs PASSPHRASE_CONFIRM
      echo >&2
      [ "$PASSPHRASE" = "$PASSPHRASE_CONFIRM" ] || { echo "error: passphrases do not match" >&2; exit 1; }
    fi

    # Derive combined key: SHA-256(vault_key_b64 + ":" + passphrase)
    # Written directly to file to avoid null-byte stripping in shell variables
    echo "Deriving combined key..."
    COMBINED_KEYFILE=$(mktemp /dev/shm/qxvault-2fa-XXXXXX)
    chmod 0600 "$COMBINED_KEYFILE"
    trap 'rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"' EXIT

    echo -n "${KEY_B64}:${PASSPHRASE}" | openssl dgst -sha256 -binary > "$COMBINED_KEYFILE"

    # Debug: show keyfile sizes
    echo "  Vault keyfile: $(wc -c < "$VAULT_KEYFILE") bytes"
    echo "  Combined keyfile: $(wc -c < "$COMBINED_KEYFILE") bytes"

    # Add new keyslot with combined key, authenticated with the vault key
    echo "Adding combined key to LUKS header (new keyslot)..."
    cryptsetup luksAddKey --key-file "$VAULT_KEYFILE" "$IMAGE" "$COMBINED_KEYFILE"

    # Verify combined key works BEFORE removing the vault slot
    echo "Verifying combined key opens LUKS..."
    if ! cryptsetup open --test-passphrase --type luks --key-file "$COMBINED_KEYFILE" "$IMAGE"; then
      echo "CRITICAL: Combined key failed to open LUKS after adding! Aborting." >&2
      echo "  Vault key is still intact. No changes were made to slot 0." >&2
      rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"
      exit 1
    fi
    echo "  Combined key verified successfully."

    # Now safely remove the vault-only keyslot
    # Find which slot the vault key occupies (usually 0)
    echo "Removing vault-only keyslot..."
    # luksKillSlot with combined key for authentication; kill slot 0
    cryptsetup luksKillSlot --key-file "$COMBINED_KEYFILE" "$IMAGE" 0

    # Final verification
    echo "Final verification..."
    if ! cryptsetup open --test-passphrase --type luks --key-file "$COMBINED_KEYFILE" "$IMAGE"; then
      echo "CRITICAL: Combined key no longer works after removing vault slot!" >&2
      echo "  This should not happen. Manual recovery may be needed." >&2
      rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"
      exit 1
    fi

    rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"

    # Save 2FA config for this VM
    cat > "${TFA_DIR}/${VMID}.conf" <<EOF
# 2FA configuration for VMID $VMID
# Generated by qxvault-2fa-enable
REQUIRE_2FA="true"
TFA_MODE="passphrase"
TFA_ENABLED_DATE="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
EOF
    chmod 0600 "${TFA_DIR}/${VMID}.conf"

    echo "2FA enabled for VM $VMID (mode: passphrase)"
    echo "  The VM can no longer auto-start. You must manually unlock via the SecureVM UI."
    echo "  The unlock requires both the vault key and your passphrase."
    ;;

  totp)
    # Use pre-generated secret if provided, otherwise generate a new one
    if [ -n "$TOTP_SECRET_OVERRIDE" ]; then
      TOTP_SECRET_B32="$TOTP_SECRET_OVERRIDE"
    else
      TOTP_SECRET_RAW=$(head -c 20 /dev/urandom)
      TOTP_SECRET_B32=$(echo -n "$TOTP_SECRET_RAW" | base32 -w0)
    fi

    # Encrypt the TOTP secret with the vault key (AES-256-CBC)
    TOTP_ENC_FILE="${TFA_DIR}/${VMID}.totp"
    echo -n "$TOTP_SECRET_B32" | openssl enc -aes-256-cbc -pbkdf2 -iter 100000 \
      -pass "pass:${KEY_B64}" -out "$TOTP_ENC_FILE"
    chmod 0600 "$TOTP_ENC_FILE"

    # Save 2FA config for this VM
    cat > "${TFA_DIR}/${VMID}.conf" <<EOF
# 2FA configuration for VMID $VMID
# Generated by qxvault-2fa-enable
REQUIRE_2FA="true"
TFA_MODE="totp"
TFA_ENABLED_DATE="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
EOF
    chmod 0600 "${TFA_DIR}/${VMID}.conf"

    # Output the provisioning URI for the user to set up their authenticator app
    TOTP_URI="otpauth://totp/SecureVM:VM${VMID}?secret=${TOTP_SECRET_B32}&issuer=SecureVM&algorithm=SHA1&digits=6&period=30"

    echo "2FA enabled for VM $VMID (mode: totp)"
    echo "  The VM can no longer auto-start. You must manually unlock via the SecureVM UI."
    echo ""
    echo "  TOTP Secret (base32): $TOTP_SECRET_B32"
    echo "  Provisioning URI:     $TOTP_URI"
    echo ""
    echo "  Add this to your authenticator app (Google Authenticator, Authy, etc.)"
    echo "  SAVE THIS SECRET — it cannot be recovered later!"
    ;;
esac

echo "Done."
