#!/bin/bash
# migrate-vm-luks — Convert an existing Proxmox VM disk to LUKS encryption
# Part of proxmox-kms-bridge (45Drives SecureVM for Proxmox)
#
# Usage: migrate-vm-luks [--dry-run] <VMID>
#
# Detects the VM's primary boot disk, creates a LUKS2 container, copies
# the disk contents into it, provisions a key in OpenBao, repoints the
# VM config, and attaches the hookscript. Saves rollback state.
set -euo pipefail

CONF_FILE="/etc/qxvault/qxvault.conf"

# ── Defaults (overridden by config) ────────────────────────────────
LUKS_IMAGE_DIR="/var/lib/vz/images"
LUKS_IMAGE_TEMPLATE="vm-{vmid}-disk-qxvault.luks"
LUKS_MAPPER_TEMPLATE="vm{vmid}_crypt"
OPENBAO_ADDR="https://127.0.0.1:8200"
OPENBAO_CACERT=""
OPENBAO_SKIP_VERIFY="false"
OPENBAO_TIMEOUT="10"

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

# ── Arguments ──────────────────────────────────────────────────────
DRY_RUN=false
VMID=""

for arg in "$@"; do
  case "$arg" in
    --dry-run) DRY_RUN=true ;;
    -h|--help)
      echo "Usage: migrate-vm-luks [--dry-run] <VMID>"
      echo "  --dry-run   Show what would happen without making changes"
      exit 0
      ;;
    *)
      if [ -z "$VMID" ]; then
        VMID="$arg"
      else
        echo "error: unexpected argument '$arg'" >&2; exit 1
      fi
      ;;
  esac
done

[ -n "$VMID" ] || { echo "Usage: migrate-vm-luks [--dry-run] <VMID>" >&2; exit 1; }
[[ "$VMID" =~ ^[0-9]+$ ]] || { echo "error: VMID must be numeric" >&2; exit 1; }

TARGET_DIR="${LUKS_IMAGE_DIR}/${VMID}"
TARGET_FILE="${TARGET_DIR}/${LUKS_IMAGE_TEMPLATE//\{vmid\}/$VMID}"
MAPPER="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$VMID}"
HOOKSCRIPT="local:snippets/qxvault-luks-hook.sh"
STATE_DIR="/root/qxvault-migrate-${VMID}"
LUKS_OVERHEAD_BYTES=$((16 * 1024 * 1024))

# ── Helpers ────────────────────────────────────────────────────────
log() { echo "[migrate-vm-luks] $*"; }
fail() { echo "[migrate-vm-luks] ERROR: $*" >&2; exit 1; }

need_cmd() {
  command -v "$1" >/dev/null 2>&1 || fail "Missing command: $1"
}

resolve_source_path() {
  local vol="$1"
  if [[ "$vol" = /* ]]; then
    printf '%s\n' "$vol"
  else
    pvesm path "$vol"
  fi
}

# ── Preflight ──────────────────────────────────────────────────────
[ "$(id -u)" -eq 0 ] || fail "Must run as root"

for c in qm pvesm qemu-img cryptsetup jq base64 awk sed grep stat truncate blockdev; do
  need_cmd "$c"
done

# Validate OpenBao connectivity before proceeding
log "Verifying OpenBao connectivity at ${OPENBAO_ADDR}..."
CURL_ARGS=(-s --connect-timeout "$OPENBAO_TIMEOUT")
if [ -n "$OPENBAO_CACERT" ]; then CURL_ARGS+=(--cacert "$OPENBAO_CACERT"); fi
if [ "$OPENBAO_SKIP_VERIFY" = "true" ]; then CURL_ARGS+=(-k); fi
HEALTH_CODE=$(curl "${CURL_ARGS[@]}" -o /dev/null -w '%{http_code}' "${OPENBAO_ADDR}/v1/sys/health" 2>/dev/null) || true
if [ "$HEALTH_CODE" != "200" ]; then
  fail "Cannot reach OpenBao at ${OPENBAO_ADDR} (HTTP ${HEALTH_CODE:-000}). Configure qxvault.conf first."
fi
log "OpenBao is reachable and healthy."

mkdir -p "$STATE_DIR"

CONFIG_MODIFIED=false

cleanup_on_error() {
  echo "[migrate-vm-luks] Migration failed. State saved in ${STATE_DIR}" >&2
  # Close mapper if we opened it
  if cryptsetup status "$MAPPER" >/dev/null 2>&1; then
    cryptsetup close "$MAPPER" 2>/dev/null || true
  fi
  # Restore original VM config if we modified it
  if [ "$CONFIG_MODIFIED" = "true" ]; then
    echo "[migrate-vm-luks] Restoring original VM config..." >&2
    qm set "$VMID" "--${BOOT_SLOT}" "$DISK_SPEC" 2>/dev/null || true
    qm set "$VMID" --delete hookscript 2>/dev/null || true
  fi
}
trap cleanup_on_error ERR

# ── Check VM state ─────────────────────────────────────────────────
if qm status "$VMID" 2>/dev/null | grep -q running; then
  fail "VM ${VMID} is running. Shut it down first."
fi

# ── Detect boot disk ──────────────────────────────────────────────
CONFIG="$(qm config "$VMID")"
echo "$CONFIG" > "${STATE_DIR}/qm-config.before.txt"

BOOT_SLOT=""
for s in scsi0 virtio0 sata0 ide0 scsi1 virtio1 sata1; do
  if echo "$CONFIG" | grep -q "^${s}: "; then
    BOOT_SLOT="$s"
    break
  fi
done
[ -n "$BOOT_SLOT" ] || fail "Could not find a disk slot (scsi0/virtio0/sata0/ide0)."

DISK_SPEC="$(echo "$CONFIG" | awk -F': ' -v s="$BOOT_SLOT" '$1==s {print $2}')"
DISK_VOLID="$(echo "$DISK_SPEC" | cut -d',' -f1)"
echo "$BOOT_SLOT" > "${STATE_DIR}/boot-slot.txt"
echo "$DISK_SPEC" > "${STATE_DIR}/disk-spec.txt"
echo "$DISK_VOLID" > "${STATE_DIR}/disk-volid.txt"

log "Detected boot disk slot: ${BOOT_SLOT}"
log "Detected disk volume: ${DISK_VOLID}"

SRC_PATH="$(resolve_source_path "$DISK_VOLID")"
echo "$SRC_PATH" > "${STATE_DIR}/source-path.txt"
[ -e "$SRC_PATH" ] || fail "Resolved source path does not exist: ${SRC_PATH}"

SRC_INFO="$(qemu-img info --output=json "$SRC_PATH")"
printf '%s\n' "$SRC_INFO" > "${STATE_DIR}/source-qemu-img-info.json"

VIRTUAL_SIZE="$(printf '%s\n' "$SRC_INFO" | jq -r '.["virtual-size"]')"
if [ -z "$VIRTUAL_SIZE" ] || [ "$VIRTUAL_SIZE" = "null" ]; then
  fail "Could not read source virtual size."
fi

SRC_FORMAT="$(printf '%s\n' "$SRC_INFO" | jq -r '.format')"
if [ -z "$SRC_FORMAT" ] || [ "$SRC_FORMAT" = "null" ]; then
  fail "Could not read source image format."
fi

case "$SRC_FORMAT" in
  raw|qcow2|qcow|vmdk|vdi|vhdx) ;;
  *) fail "Unsupported source format: ${SRC_FORMAT}" ;;
esac

TARGET_SIZE=$((VIRTUAL_SIZE + LUKS_OVERHEAD_BYTES))
echo "$VIRTUAL_SIZE" > "${STATE_DIR}/source-virtual-size.txt"
echo "$TARGET_SIZE" > "${STATE_DIR}/target-file-size.txt"
echo "$SRC_FORMAT" > "${STATE_DIR}/source-format.txt"

log "Source path: ${SRC_PATH}"
log "Source format: ${SRC_FORMAT}"
log "Virtual size: ${VIRTUAL_SIZE} bytes"
log "Target file size: ${TARGET_SIZE} bytes"
log "Target LUKS file: ${TARGET_FILE}"

# ── Dry run ────────────────────────────────────────────────────────
if [ "$DRY_RUN" = "true" ]; then
  log "DRY RUN — no changes made. Summary:"
  log "  VM ${VMID}: ${BOOT_SLOT} = ${DISK_VOLID}"
  log "  Source: ${SRC_PATH} (${SRC_FORMAT}, ${VIRTUAL_SIZE} bytes)"
  log "  Target: ${TARGET_FILE} (${TARGET_SIZE} bytes)"
  log "  Hookscript: ${HOOKSCRIPT}"
  log "  Key will be provisioned in OpenBao transit engine (key name: ${VMID})"
  exit 0
fi

# ── Create target ─────────────────────────────────────────────────
install -d -m 700 "$TARGET_DIR"

# ── Operation lock file (for UI progress tracking) ───────────────
OP_LOCK="/run/qxvault-op-${VMID}.json"
echo "{\"op\":\"encrypt\",\"vmid\":\"${VMID}\",\"startTime\":$(date +%s)}" > "$OP_LOCK"
trap 'rm -f "$OP_LOCK"' EXIT

if [ -e "$TARGET_FILE" ]; then
  fail "Target file already exists: ${TARGET_FILE}"
fi

# ── Disk space check ──────────────────────────────────────────────
AVAIL_BYTES=$(df --output=avail -B1 "$TARGET_DIR" | tail -1 | tr -d ' ')
REQUIRED_BYTES=$(( TARGET_SIZE + TARGET_SIZE / 10 ))  # 10% safety margin
if [ "$REQUIRED_BYTES" -gt "$AVAIL_BYTES" ]; then
  fail "Insufficient disk space: need $(numfmt --to=iec "$REQUIRED_BYTES") (incl. 10% margin), have $(numfmt --to=iec "$AVAIL_BYTES") on $(df --output=target "$TARGET_DIR" | tail -1)"
fi


# Verify key retrieval works
KEY_BYTES="$(/usr/libexec/vault-dmkey get-key --name "$VMID" | base64 -d | wc -c | tr -d ' ')"
[ "$KEY_BYTES" = "32" ] || fail "Vault key for VM ${VMID} is not 32 bytes after decode."

log "Creating target file: ${TARGET_FILE}"
truncate -s "$TARGET_SIZE" "$TARGET_FILE"
chmod 600 "$TARGET_FILE"

log "Formatting target as LUKS2"
/usr/libexec/vault-dmkey get-key --name "$VMID" \
  | base64 -d \
  | cryptsetup luksFormat --type luks2 --key-file=- "$TARGET_FILE"

log "Opening mapper: ${MAPPER}"
/usr/libexec/vault-dmkey get-key --name "$VMID" \
  | base64 -d \
  | cryptsetup open --type luks --key-file=- "$TARGET_FILE" "$MAPPER"

MAPPER_SIZE="$(blockdev --getsize64 "/dev/mapper/${MAPPER}")"
echo "$MAPPER_SIZE" > "${STATE_DIR}/mapper-size.txt"
log "Mapped device size: ${MAPPER_SIZE} bytes"

if [ "$MAPPER_SIZE" -lt "$VIRTUAL_SIZE" ]; then
  cryptsetup close "$MAPPER"
  fail "Mapped device is smaller than source image."
fi

log "Copying source disk into /dev/mapper/${MAPPER} ..."
qemu-img convert -t none -p -n -f "$SRC_FORMAT" -O raw "$SRC_PATH" "/dev/mapper/${MAPPER}"

# Verify the written data is readable and has correct size
WRITTEN_SIZE="$(blockdev --getsize64 "/dev/mapper/${MAPPER}")"
if [ "$WRITTEN_SIZE" -lt "$VIRTUAL_SIZE" ]; then
  fail "Post-copy verification failed: written ${WRITTEN_SIZE} bytes but expected at least ${VIRTUAL_SIZE}"
fi
log "Post-copy verification passed (${WRITTEN_SIZE} bytes on mapper)"

log "Verifying Proxmox snippets storage"
install -d -m 755 /var/lib/vz/snippets

if [ ! -f /var/lib/vz/snippets/qxvault-luks-hook.sh ]; then
  fail "Missing /var/lib/vz/snippets/qxvault-luks-hook.sh — is proxmox-kms-bridge installed?"
fi

log "Attaching hookscript"
CONFIG_MODIFIED=true
qm set "$VMID" --hookscript "$HOOKSCRIPT"

log "Switching VM disk on ${BOOT_SLOT} to /dev/mapper/${MAPPER}"
qm set "$VMID" "--${BOOT_SLOT}" "/dev/mapper/${MAPPER}"

NEWCFG="$(qm config "$VMID")"
echo "$NEWCFG" > "${STATE_DIR}/qm-config.after.txt"

log "=========================================="
log "  Migration complete for VMID ${VMID}"
log "=========================================="
log "  LUKS file:      ${TARGET_FILE}"
log "  Mapper device:  /dev/mapper/${MAPPER}"
log "  Hookscript:     ${HOOKSCRIPT}"
log "  Key stored in:  OpenBao transit engine (key name: ${VMID})"
log "  Rollback state: ${STATE_DIR}/"
log ""
log "  Start the VM:   qm start ${VMID}"
log "=========================================="