#!/usr/bin/env bash
set -euo pipefail

# luks-img.sh — manage LUKS-encrypted .img files (loop, open, mount)
# Now with:
#   launch-vbox  — attach unlocked mapper to a VirtualBox VM (raw VMDK wrapper)
#   detach-vbox  — detach & remove the wrapper
#   attach-virt  — attach unlocked mapper to a libvirt/QEMU VM
#   detach-virt  — detach it again
#
# Requires: cryptsetup, losetup, mount, umount, findmnt, blkid
# Optional: VBoxManage (VirtualBox), virsh (libvirt)

die() { echo "ERROR: $*" >&2; exit 1; }
need() { command -v "$1" >/dev/null 2>&1 || die "Missing dependency: $1"; }
banner() { echo -e "\n== $* ==\n"; }

SUDO=${SUDO:-}
as_root() {
  if [[ $EUID -ne 0 ]]; then
    if [[ -z "${SUDO}" ]]; then
      exec sudo SUDO=1 -- "$0" "$@"
    else
      die "Escalation loop detected; aborting."
    fi
  fi
}

usage() {
  cat <<'EOF'
luks-img.sh — Create/open/mount/close LUKS-encrypted disk images (.img)

USAGE:
  luks-img.sh create   --file FILE.img --size 4G [--fs ext4] [--name NAME] [--mount /mnt/foo]
                       [--cipher aes-xts-plain64 --key-size 512 --hash sha512 --pbkdf argon2id]
                       [--keyfile /path/key] [--sparse] [--no-mkfs]
  luks-img.sh open     --file FILE.img [--name NAME] [--mount /mnt/foo] [--keyfile /path/key]
  luks-img.sh close    --file FILE.img [--name NAME]
  luks-img.sh status   --name NAME
  luks-img.sh dump     --file FILE.img
  luks-img.sh add-key  --file FILE.img --keyfile /path/new_key [--existing-keyfile /path/key]
  luks-img.sh header-backup  --file FILE.img --out FILE.header
  luks-img.sh header-restore --file FILE.img --in  FILE.header

  # NEW: VirtualBox & virt-manager helpers
  luks-img.sh launch-vbox --file FILE.img --vm "VM Name" [--name NAME] [--keyfile /path/key] \
                          [--port 0] [--controller "SATA Controller"] [--start]
  luks-img.sh detach-vbox --vm "VM Name" [--port 0] [--controller "SATA Controller"] [--wrapper PATH.vmdk]

  luks-img.sh attach-virt --file FILE.img --vm myvm [--name NAME] [--keyfile /path/key] \
                          [--target vdb] [--bus virtio] [--persistent]
  luks-img.sh detach-virt --vm myvm [--target vdb] [--persistent]

NOTES:
- NAME defaults to basename(FILE) without extension (secure.img -> secure).
- 'launch-vbox' creates a raw VMDK wrapper pointing at /dev/mapper/NAME, attaches it, and can --start the VM.
- 'detach-vbox' removes that attachment and deletes the wrapper file (path is shown by launch-vbox).
- 'attach-virt' uses virsh attach-disk to add /dev/mapper/NAME to a libvirt VM (virt-manager will show it).
- 'detach-virt' removes that attachment.

EXAMPLES:
  sudo ./luks-img.sh create --file secure.img --size 8G --mount /mnt/secure
  sudo ./luks-img.sh open   --file secure.img --mount /mnt/secure
  sudo ./luks-img.sh launch-vbox --file secure.img --vm "Ubuntu VM" --start
  sudo ./luks-img.sh attach-virt --file secure.img --vm myvm --target vdb --persistent
  sudo ./luks-img.sh detach-vbox --vm "Ubuntu VM"
  sudo ./luks-img.sh detach-virt --vm myvm --target vdb --persistent
EOF
  exit 1
}

# ---------- tiny arg helpers ----------
expand_kv() { [[ "$1" == --*=* ]] && { k="${1%%=*}"; v="${1#*=}"; set -- "$k" "$v" "${@:2}"; echo "$@"; return 0; }; echo "$@"; }

# ---------- core helpers ----------
find_loop_by_file() { losetup -j "$1" | awk -F: '{print $1}' | head -n1; }
ensure_loop() { local f="$1"; local l; l=$(find_loop_by_file "$f" || true); [[ -z "$l" ]] && l=$(losetup --find --show "$f"); echo "$l"; }
mapper_path() { echo "/dev/mapper/$1"; }
is_mapped() { [[ -e "$(mapper_path "$1")" ]]; }

mount_if_requested() { local dev="$1" mnt="${2:-}"; [[ -z "$mnt" ]] && return 0; mkdir -p "$mnt"; mount "$dev" "$mnt"; echo "$mnt"; }
umount_if_mounted() {
  local dev="$1"
  while read -r t; do [[ -n "$t" ]] && umount "$t"; done < <(findmnt -n -o TARGET "$dev" 2>/dev/null || true)
}
mkfs_for() {
  local fs="$1" dev="$2"
  case "$fs" in
    ext4) mkfs.ext4 -F "$dev" ;;
    xfs)  mkfs.xfs -f "$dev" ;;
    btrfs) mkfs.btrfs -f "$dev" ;;
    *) die "Unsupported/unknown fs: $fs" ;;
  esac
}

# ---------- actions (create/open/close/etc.) ----------
do_create() {
  local file="" size="" fs="ext4" name="" mount_dir=""
  local cipher="aes-xts-plain64" key_size="512" hash="sha512" pbkdf="argon2id"
  local sparse=0 no_mkfs=0 keyfile=""

  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --file) shift; file="$1" ;;
      --size) shift; size="$1" ;;
      --fs) shift; fs="$1" ;;
      --name) shift; name="$1" ;;
      --mount) shift; mount_dir="$1" ;;
      --cipher) shift; cipher="$1" ;;
      --key-size) shift; key_size="$1" ;;
      --hash) shift; hash="$1" ;;
      --pbkdf) shift; pbkdf="$1" ;;
      --keyfile) shift; keyfile="$1" ;;
      --sparse) sparse=1 ;;
      --no-mkfs) no_mkfs=1 ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done

  [[ -z "$file" || -z "$size" ]] && usage
  [[ -z "$name" ]] && name="$(basename "${file%.*}")"

  need cryptsetup; need losetup; need mount; need findmnt; as_root "$@"

  banner "Creating image file ($size) at $file"
  if (( sparse )); then
    fallocate -l "$size" "$file" 2>/dev/null || dd if=/dev/zero of="$file" bs=1 count=0 seek="$size"
  else
    fallocate -l "$size" "$file" 2>/dev/null || dd if=/dev/zero of="$file" bs=1M count="${size%G}" seek=0
  fi

  banner "Attaching loop device"
  local loop; loop=$(ensure_loop "$file"); echo "Loop: $loop"

  banner "Initializing LUKS2"
  local luks_args=(--type luks2 --cipher "$cipher" --key-size "$key_size" --hash "$hash" --pbkdf "$pbkdf")
  if [[ -n "$keyfile" ]]; then
    [[ -f "$keyfile" ]] || die "Keyfile not found: $keyfile"
    cryptsetup luksFormat "${luks_args[@]}" --batch-mode --key-file "$keyfile" "$loop"
  else
    cryptsetup luksFormat "${luks_args[@]}" "$loop"
  fi

  banner "Opening mapper $name"
  if [[ -n "$keyfile" ]]; then cryptsetup luksOpen "$loop" "$name" --key-file "$keyfile"
  else cryptsetup luksOpen "$loop" "$name"; fi

  local dev="/dev/mapper/$name"
  if (( ! no_mkfs )); then banner "mkfs.$fs on $dev"; mkfs_for "$fs" "$dev"; fi
  if [[ -n "$mount_dir" ]]; then banner "Mounting at $mount_dir"; mount_if_requested "$dev" "$mount_dir" >/dev/null; fi

  echo "DONE. Loop=$loop Mapper=$dev"
}

do_open() {
  local file="" name="" mount_dir="" keyfile=""
  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --file) shift; file="$1" ;;
      --name) shift; name="$1" ;;
      --mount) shift; mount_dir="$1" ;;
      --keyfile) shift; keyfile="$1" ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done
  [[ -z "$file" ]] && usage
  [[ -z "$name" ]] && name="$(basename "${file%.*}")"
  need cryptsetup; need losetup; need mount; need findmnt; as_root "$@"

  banner "Attach loop"
  local loop; loop=$(ensure_loop "$file"); echo "Loop: $loop"
  banner "Open LUKS -> $name"
  if [[ -n "$keyfile" ]]; then cryptsetup luksOpen "$loop" "$name" --key-file "$keyfile"
  else cryptsetup luksOpen "$loop" "$name"; fi

  local dev="/dev/mapper/$name"
  [[ -n "$mount_dir" ]] && { banner "Mounting $dev at $mount_dir"; mount_if_requested "$dev" "$mount_dir" >/dev/null; }
  echo "Mapper: $dev"
}

do_close() {
  local file="" name=""
  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --file) shift; file="$1" ;;
      --name) shift; name="$1" ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done
  [[ -z "$file" ]] && usage
  [[ -z "$name" ]] && name="$(basename "${file%.*}")"
  need cryptsetup; need losetup; need findmnt; as_root "$@"

  local loop; loop=$(find_loop_by_file "$file" || true)
  local dev="/dev/mapper/$name"

  if is_mapped "$name"; then
    banner "Unmount mounts for $dev (if any)"
    umount_if_mounted "$dev" || true
    banner "cryptsetup luksClose $name"
    cryptsetup luksClose "$name"
  else
    echo "Mapper not open: $name"
  fi

  if [[ -n "$loop" ]]; then banner "Detach loop $loop"; losetup -d "$loop"; else echo "No loop device for $file"; fi
  echo "Closed."
}

do_status() { local name=""; while (( $# )); do case "$1" in --name) shift; name="$1";; *) usage;; esac; shift||true; done; [[ -z "$name" ]]&&usage; need cryptsetup; cryptsetup status "$name"||true; findmnt "/dev/mapper/$name"||true; }
do_dump() { local file=""; while (( $# )); do case "$1" in --file) shift; file="$1";; *) usage;; esac; shift||true; done; [[ -z "$file" ]]&&usage; need cryptsetup; cryptsetup luksDump "$file"; }
do_add_key() { local file="" keyfile="" existing=""; while (( $# )); do case "$1" in --file) shift; file="$1";; --keyfile) shift; keyfile="$1";; --existing-keyfile) shift; existing="$1";; *) usage;; esac; shift||true; done; [[ -z "$file"||-z "$keyfile" ]]&&usage; need cryptsetup; as_root "$@"; [[ -n "$existing" ]] && cryptsetup luksAddKey "$file" --key-file "$existing" "$keyfile" || cryptsetup luksAddKey "$file" "$keyfile"; echo "Key added."; }
do_header_backup() { local file="" out=""; while (( $# )); do case "$1" in --file) shift; file="$1";; --out) shift; out="$1";; *) usage;; esac; shift||true; done; [[ -z "$file"||-z "$out" ]]&&usage; need cryptsetup; as_root "$@"; cryptsetup luksHeaderBackup "$file" --header-backup-file "$out"; echo "Header -> $out"; }
do_header_restore() { local file="" in=""; while (( $# )); do case "$1" in --file) shift; file="$1";; --in) shift; in="$1";; *) usage;; esac; shift||true; done; [[ -z "$file"||-z "$in" ]]&&usage; need cryptsetup; as_root "$@"; cryptsetup luksHeaderRestore "$file" --header-backup-file "$in"; echo "Header restored from $in"; }

# ---------- VirtualBox ----------
do_launch_vbox() {
  need VBoxManage
  local file="" name="" keyfile="" vm="" controller="SATA Controller" port="0" start_vm=0
  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --file) shift; file="$1" ;;
      --name) shift; name="$1" ;;
      --keyfile) shift; keyfile="$1" ;;
      --vm) shift; vm="$1" ;;
      --controller) shift; controller="$1" ;;
      --port) shift; port="$1" ;;
      --start) start_vm=1 ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done
  [[ -z "$file" || -z "$vm" ]] && usage
  [[ -z "$name" ]] && name="$(basename "${file%.*}")"

  # Open (host-side decrypt) to /dev/mapper/NAME
  "$0" open --file "$file" --name "$name" ${keyfile:+--keyfile "$keyfile"} >/dev/null

  local mapper; mapper=$(mapper_path "$name")
  [[ -b "$mapper" ]] || die "Mapper not found: $mapper"

  # Create a raw VMDK wrapper pointing to the mapper
  local cache_dir="${XDG_CACHE_HOME:-$HOME/.cache}/luks-img"
  mkdir -p "$cache_dir"
  local wrapper="$cache_dir/${name}.vmdk"

  banner "Creating raw VMDK wrapper: $wrapper -> $mapper"
  # VBox needs read perms on the mapper; run VirtualBox as same user who owns it.
  VBoxManage internalcommands createrawvmdk -filename "$wrapper" -rawdisk "$mapper" >/dev/null

  banner "Attaching to VM: $vm (controller=$controller, port=$port)"
  VBoxManage storageattach "$vm" \
    --storagectl "$controller" \
    --port "$port" --device 0 --type hdd --medium "$wrapper"

  echo "VMDK wrapper: $wrapper"
  if (( start_vm )); then
    banner "Starting VM: $vm"
    VBoxManage startvm "$vm" --type gui
  else
    echo "Attached. Start the VM when ready."
  fi
  echo "NOTE: Use 'detach-vbox' to safely remove and delete the wrapper."
}

do_detach_vbox() {
  need VBoxManage
  local vm="" controller="SATA Controller" port="0" wrapper=""
  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --vm) shift; vm="$1" ;;
      --controller) shift; controller="$1" ;;
      --port) shift; port="$1" ;;
      --wrapper) shift; wrapper="$1" ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done
  [[ -z "$vm" ]] && usage

  banner "Detaching disk from VM: $vm (controller=$controller, port=$port)"
  VBoxManage storageattach "$vm" --storagectl "$controller" --port "$port" --device 0 --medium none || true

  if [[ -n "$wrapper" && -f "$wrapper" ]]; then
    banner "Deleting wrapper $wrapper"
    rm -f -- "$wrapper"
  else
    echo "If you used launch-vbox, the wrapper is in \$XDG_CACHE_HOME/luks-img/<name>.vmdk"
  fi

  echo "Done. (This does NOT close your LUKS mapper—run 'close' if needed.)"
}

# ---------- libvirt / virt-manager ----------
do_attach_virt() {
  need virsh
  local file="" name="" keyfile="" vm="" target="vdb" bus="virtio" persistent=0
  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --file) shift; file="$1" ;;
      --name) shift; name="$1" ;;
      --keyfile) shift; keyfile="$1" ;;
      --vm) shift; vm="$1" ;;
      --target) shift; target="$1" ;;
      --bus) shift; bus="$1" ;;
      --persistent) persistent=1 ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done
  [[ -z "$file" || -z "$vm" ]] && usage
  [[ -z "$name" ]] && name="$(basename "${file%.*}")"

  "$0" open --file "$file" --name "$name" ${keyfile:+--keyfile "$keyfile"} >/dev/null

  local mapper; mapper=$(mapper_path "$name")
  [[ -b "$mapper" ]] || die "Mapper not found: $mapper"

  banner "Attaching $mapper to libvirt VM '$vm' as $target (bus=$bus)"
  if (( persistent )); then
    virsh attach-disk "$vm" "$mapper" "$target" --targetbus "$bus" --persistent --driver qemu || die "virsh attach-disk failed"
  else
    virsh attach-disk "$vm" "$mapper" "$target" --targetbus "$bus" --driver qemu || die "virsh attach-disk failed"
  fi

  echo "Attached. (Virt-manager will show the new disk.)"
}

do_detach_virt() {
  need virsh
  local vm="" target="vdb" persistent=0
  while (( $# )); do
    set -- $(expand_kv "$@")
    case "$1" in
      --vm) shift; vm="$1" ;;
      --target) shift; target="$1" ;;
      --persistent) persistent=1 ;;
      --help|-h) usage ;;
      *) die "Unknown arg: $1" ;;
    esac; shift || true
  done
  [[ -z "$vm" ]] && usage

  banner "Detaching $target from libvirt VM '$vm'"
  if (( persistent )); then
    virsh detach-disk "$vm" "$target" --persistent || die "virsh detach-disk failed"
  else
    virsh detach-disk "$vm" "$target" || die "virsh detach-disk failed"
  fi

  echo "Detached. (This does NOT close your LUKS mapper—run 'close' if needed.)"
}

# ---------- main ----------
[[ $# -lt 1 ]] && usage
sub="$1"; shift || true

case "$sub" in
  create)            do_create "$@" ;;
  open)              do_open "$@" ;;
  close)             do_close "$@" ;;
  status)            do_status "$@" ;;
  dump)              do_dump "$@" ;;
  add-key)           do_add_key "$@" ;;
  header-backup)     do_header_backup "$@" ;;
  header-restore)    do_header_restore "$@" ;;
  launch-vbox)       do_launch_vbox "$@" ;;
  detach-vbox)       do_detach_vbox "$@" ;;
  attach-virt)       do_attach_virt "$@" ;;
  detach-virt)       do_detach_virt "$@" ;;
  -h|--help|help)    usage ;;
  *)                 die "Unknown subcommand: $sub (use --help)" ;;
esac