#!/bin/bash
# Boots a sealed Omarchy image (vm-seal) as a throwaway machine.
#
#   vm-boot <image>                 boot, print how to log in, run until Ctrl-C
#   vm-boot <image> <command...>    boot, run the command over SSH, power off;
#                                   exits with the command's status
#
# <image> is a downloaded disk, omarchy-<...>.qcow2, with vmlinuz, initrd and
# cmdline in the same directory as this script, as in a published "unpacked"
# folder. Or it is a directory as vm-seal leaves it, holding base.qcow2 and
# those three files.
#
# Nothing is written to the image: the machine runs on a copy-on-write overlay
# that is deleted on exit. Each boot gets a new SSH key and a new machine ID;
# the image itself holds no key and its SSH takes no password.
#
# Needs qemu-system-x86_64, qemu-img, ssh and a usable /dev/kvm.
# Env: SMP (4), MEM (4G), SSH_PORT (random), VM_UUID (random), APPEND (more
#      kernel arguments), QEMU_EXTRA (more qemu arguments, split on spaces),
#      VM_STATE_DIR (/var/tmp; holds the overlay, so it needs room to grow).
set -euo pipefail
[[ $# -ge 1 ]] || { sed -n '2,21p' "$0" >&2; exit 2; }
image=$(realpath "$1"); shift
if [[ -d $image ]]; then
  base=$image disk=$image/base.qcow2
else
  base=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) disk=$image
fi
[[ -s $disk ]] || { echo "vm-boot: there is no disk image at $disk" >&2; exit 1; }
for f in vmlinuz initrd cmdline; do
  [[ -s $base/$f ]] || { echo "vm-boot: $base/$f is missing" >&2; exit 1; }
done
[[ -r /dev/kvm && -w /dev/kvm ]] || { echo "vm-boot: /dev/kvm is not usable" >&2; exit 1; }

state=$(mktemp -d "${VM_STATE_DIR:-/var/tmp}/vm.XXXXXX")
cleanup() {
  local status=$? pid
  if [[ -f $state/qemu.pid ]]; then
    pid=$(cat "$state/qemu.pid")
    kill "$pid" 2>/dev/null || true
    for _ in $(seq 1 50); do kill -0 "$pid" 2>/dev/null || break; sleep 0.1; done
  fi
  rm -rf "$state"
  exit "$status"
}
trap cleanup EXIT
trap 'exit 130' INT TERM
port=${SSH_PORT:-$(( 20000 + RANDOM % 10000 ))}
uuid=${VM_UUID:-$(cat /proc/sys/kernel/random/uuid)}

# A key made for this boot only, handed in through the firmware tables as a
# systemd credential that writes authorized_keys.
ssh-keygen -q -t ed25519 -N '' -C vm-boot -f "$state/id_ed25519"
keys=$(printf 'd /home/omarchy/.ssh 0700 omarchy omarchy -\nf+ /home/omarchy/.ssh/authorized_keys 0600 omarchy omarchy - %s\n' \
  "$(cat "$state/id_ed25519.pub")" | base64 -w0)
ssh_opts=(-i "$state/id_ed25519" -p "$port" -o BatchMode=yes -o IdentitiesOnly=yes
  -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR)

qemu-img create -qf qcow2 -b "$disk" -F qcow2 "$state/run.qcow2"
t0=$(date +%s%N)
# shellcheck disable=SC2086
qemu-system-x86_64 -cpu host -enable-kvm -machine q35,accel=kvm -smp "${SMP:-4}" -m "${MEM:-4G}" \
  -kernel "$base/vmlinuz" -initrd "$base/initrd" -append "$(cat "$base/cmdline")${APPEND:+ $APPEND}" \
  -uuid "$uuid" \
  -drive "file=$state/run.qcow2,format=qcow2,if=none,id=d0" -device virtio-blk-pci,drive=d0 \
  -smbios "type=11,value=io.systemd.credential.binary:tmpfiles.extra=$keys" \
  -device virtio-vga -device virtio-rng-pci -display none \
  -netdev "user,id=n0,hostfwd=tcp:127.0.0.1:$port-:22" -device virtio-net-pci,netdev=n0 \
  -qmp "unix:$state/qmp.sock,server,nowait" -pidfile "$state/qemu.pid" -daemonize ${QEMU_EXTRA:-}

# </dev/null: the probe must not eat a script meant for the command below.
until ssh "${ssh_opts[@]}" -o ConnectTimeout=1 omarchy@127.0.0.1 true </dev/null 2>/dev/null; do
  kill -0 "$(cat "$state/qemu.pid")" 2>/dev/null || { echo "vm-boot: qemu exited before ssh came up" >&2; exit 1; }
  (( ($(date +%s%N) - t0) / 1000000 < 300000 )) || { echo "vm-boot: no ssh after 300s" >&2; exit 1; }
  sleep 0.2
done
echo "vm-boot: ssh up after $(( ($(date +%s%N) - t0) / 1000000 )) ms, machine ID ${uuid//-/}" >&2

if [[ $# -gt 0 ]]; then
  # Stdin passes through, so `vm-boot dir bash -s <script` works.
  ssh "${ssh_opts[@]}" omarchy@127.0.0.1 "$@"
else
  echo "ssh ${ssh_opts[*]} omarchy@127.0.0.1"
  echo "Desktop and sudo password: omarchy. QMP socket: $state/qmp.sock. Ctrl-C powers it off."
  while kill -0 "$(cat "$state/qemu.pid")" 2>/dev/null; do sleep 1; done
fi