#!/bin/bash
# qxvault-2fa-disable — Disable 2FA (second-pass unlock) for a SecureVM-encrypted VM
# Part of proxmox-kms-bridge (45Drives SecureVM for Proxmox)
#
# Usage: qxvault-2fa-disable <vmid> [passphrase]
#
# In passphrase mode:
#   - Re-derives the combined key from vault_key + passphrase
#   - Adds the vault-only key back as a LUKS keyslot
#   - Removes the combined key keyslot
#   - Restores automatic unlock on VM start
#
# In TOTP mode:
#   - Removes the TOTP secret file
#   - Removes the 2FA config (hookscript resumes auto-unlock)
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:-}"
PASSPHRASE="${2:-}"

[ -n "$VMID" ] || { echo "usage: qxvault-2fa-disable <vmid> [passphrase]" >&2; exit 1; }
[[ "$VMID" =~ ^[0-9]+$ ]] || { echo "error: VMID must be numeric" >&2; exit 1; }

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

# Check 2FA is enabled
TFA_CONF="${TFA_DIR}/${VMID}.conf"
[ -f "$TFA_CONF" ] || { echo "error: 2FA is not enabled for VMID $VMID" >&2; exit 1; }

# Read current 2FA mode
TFA_MODE=""
# shellcheck source=/dev/null
. "$TFA_CONF"

[ -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

# ── Retrieve 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

# ── Mode-specific disable ─────────────────────────────────────────
case "$TFA_MODE" in
  passphrase)
    # Read passphrase if not passed
    if [ -z "$PASSPHRASE" ]; then
      echo "Enter current 2FA passphrase for VM $VMID:" >&2
      read -rs PASSPHRASE
      echo >&2
      [ -n "$PASSPHRASE" ] || { echo "error: passphrase cannot be empty" >&2; exit 1; }
    fi

    # Re-derive the combined key (write directly to file to avoid null-byte issues)
    echo "Deriving combined key..."
    COMBINED_KEYFILE=$(mktemp /dev/shm/qxvault-2fa-XXXXXX)
    chmod 0600 "$COMBINED_KEYFILE"
    trap 'rm -f "$COMBINED_KEYFILE"' EXIT
    echo -n "${KEY_B64}:${PASSPHRASE}" | openssl dgst -sha256 -binary > "$COMBINED_KEYFILE"

    echo "Verifying combined key against LUKS..."
    if ! cryptsetup open --test-passphrase --type luks --key-file "$COMBINED_KEYFILE" "$IMAGE"; then
      echo "error: passphrase is incorrect (combined key does not match LUKS)" >&2
      rm -f "$COMBINED_KEYFILE"
      exit 1
    fi

    # Add vault-only key back
    echo "Re-adding vault-only key to LUKS..."
    trap 'rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"' EXIT

    cryptsetup luksAddKey --key-file "$COMBINED_KEYFILE" "$IMAGE" "$VAULT_KEYFILE"

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

    # Now safely remove the combined keyslot using luksRemoveKey
    # luksRemoveKey removes the slot matching the provided key content
    echo "Removing combined key from LUKS..."
    cryptsetup luksRemoveKey "$IMAGE" --key-file "$COMBINED_KEYFILE" 2>/dev/null \
      || echo "  (combined slot already removed or cleaned up)"

    # Final verification
    echo "Final verification..."
    if ! cryptsetup open --test-passphrase --type luks --key-file "$VAULT_KEYFILE" "$IMAGE"; then
      echo "CRITICAL: Vault key verification failed after removing combined slot!" >&2
      rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"
      exit 1
    fi

    rm -f "$COMBINED_KEYFILE" "$VAULT_KEYFILE"
    ;;

  totp)
    # TOTP mode: just remove the TOTP secret and config
    echo "Removing TOTP secret..."
    rm -f "${TFA_DIR}/${VMID}.totp"

    # Verify the vault key still works (it should — TOTP doesn't rekey LUKS)
    echo "Verifying vault key against LUKS..."
    if ! cryptsetup open --test-passphrase --type luks --key-file "$VAULT_KEYFILE" "$IMAGE"; then
      echo "warning: vault key does not match LUKS. LUKS may have been rekeyed manually." >&2
    fi
    ;;

  *)
    echo "error: unknown TFA_MODE '$TFA_MODE' in ${TFA_CONF}" >&2
    exit 1
    ;;
esac

# ── Remove 2FA config ──────────────────────────────────────────────
rm -f "$TFA_CONF"
echo "2FA disabled for VM $VMID"
echo "  The VM will now auto-unlock on start (hookscript resumes normal behavior)."
echo "Done."
