#!/usr/bin/env bash
# Build the VM's virtio disk image carrying the device's canonical 12-partition
# A/B layout (qemu/blkdevparts.conf: the same string U-Boot and Linux parse on
# hardware; there is no MBR/GPT). Every partition is placed at the exact offset
# the cmdline string declares; rootfs_a/rootfs_b/oem_a/oem_b/userdata get ext4,
# the boot-chain partitions (env/idblock/uboot/misc/boot_a/boot_b/recovery)
# stay zeroed: the VM enters at -kernel and never reads them.
#
# Built entirely UNPRIVILEGED: per-partition mkfs.ext4 -d (no loop mounts, no
# sudo), then dd'd into a sparse raw image.
#
# Usage: mkimage.sh [--portal-url URL] [--state KEY=VALUE]... [--fw-version V]
#                   [--rootfs-image PATH] [--oem-image PATH]
# --rootfs-image/--oem-image place REAL device images (raw ext4, e.g. a
# flare-edge build's rootfs.img/oem.img matched pair) into slot A instead of
# the busybox skeleton; slot B keeps the skeleton as a known-good fallback.
# Env:
#   BUSYBOX    path to a local busybox binary (skips the download; still verified)
#   OUT        output dir (default: qemu/out); image at $OUT/disk.img
#   SEED_DIR   a directory of pre-built userdata/warden files (one file per
#              settings key, e.g. flare-edge tools/seed-fixtures.py's output)
#              copied in VERBATIM, after --portal-url/--state. Lets a caller
#              seed an arbitrary key set -- a flow spec's setup.seed can name
#              anything settings.c reads -- without growing --state into a
#              multi-value flag; last one written wins, same as --state.
set -euo pipefail

QEMU_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=qemu/lib.sh disable=SC1091
. "$QEMU_DIR/lib.sh"
# shellcheck source=qemu/blkdevparts.conf disable=SC1091
. "$QEMU_DIR/blkdevparts.conf"
OUT="${OUT:-$QEMU_DIR/out}"
# mkfs.ext4 lives in sbin, which user shells on Debian don't have on PATH.
PATH="$PATH:/usr/sbin:/sbin"

PORTAL_URL=""
STATE_KV=()
FW_VERSION="0.0.1"
ROOTFS_IMAGE=""
OEM_IMAGE=""
while [ $# -gt 0 ]; do
  case "$1" in
    --portal-url) PORTAL_URL="${2:?--portal-url needs a value}"; shift 2 ;;
    --state)
      case "${2:?--state needs KEY=VALUE}" in
        *=*) ;;
        *) echo "FATAL: --state needs KEY=VALUE, got '$2'" >&2; exit 1 ;;
      esac
      case "${2%%=*}" in
        .|..)
          echo "FATAL: --state key cannot be '.' or '..'" >&2
          exit 1 ;;
        *[!A-Za-z0-9_.]*|'')
          echo "FATAL: --state key '${2%%=*}' must match [A-Za-z0-9_.]+ (it becomes a filename)" >&2
          exit 1 ;;
      esac
      STATE_KV+=("$2"); shift 2 ;;
    --fw-version) FW_VERSION="${2:?--fw-version needs a value}"; shift 2 ;;
    --rootfs-image) ROOTFS_IMAGE="${2:?--rootfs-image needs a path}"; shift 2 ;;
    --oem-image)    OEM_IMAGE="${2:?--oem-image needs a path}"; shift 2 ;;
    *) echo "FATAL: unknown argument '$1' (usage: mkimage.sh [--portal-url URL] [--state KEY=VALUE]... [--fw-version V])" >&2; exit 1 ;;
  esac
done

qemu_get_busybox

SCRATCH="$(mktemp -d "${TMPDIR:-/tmp}/warden-qemu-image.XXXXXX")"
trap 'rm -rf "$SCRATCH"' EXIT

# Stage the rootfs tree once (skeleton + payload into /usr/bin), reused for
# both slots so A and B start byte-identical, like a factory flash.
ROOT="$SCRATCH/root"
qemu_stage_rootfs "$ROOT"
for p in "$QEMU_DIR"/payload/*; do
  [ -f "$p" ] || continue
  case "$(basename "$p")" in README.md) continue ;; esac
  install -m 0755 "$p" "$ROOT/usr/bin/$(basename "$p")"
done

# Firmware version stamp: same path the device build writes; flared reads its
# running version here (downgrade rules key off it).
printf '%s\n' "$FW_VERSION" > "$ROOT/etc/warden-firmware-version"

# Seed persistent state (flared: one file per key under /userdata/warden).
# Newline-terminated, matching how flare-edge's fw-e2e-test.sh seeds the store.
UDATA="$SCRATCH/userdata"
mkdir -p "$UDATA/warden"
[ -n "$PORTAL_URL" ] && printf '%s\n' "$PORTAL_URL" > "$UDATA/warden/flare.url"
for kv in ${STATE_KV[@]+"${STATE_KV[@]}"}; do
  printf '%s\n' "${kv#*=}" > "$UDATA/warden/${kv%%=*}"
done

# SEED_DIR: a pre-built set of settings files (see the Env note above), copied
# in whole rather than re-parsed here -- seed-fixtures.py already wrote them
# in the exact format settings.c reads (filename = key, mode 0600 for a
# secret), so re-deriving that here would be a second place to keep in sync
# with settings.c. Applied after --state so a seeded file can override a
# same-named --state value; `cp -a` preserves the 0600 on a secret entry.
if [ -n "${SEED_DIR:-}" ]; then
  [ -d "$SEED_DIR" ] || { echo "FATAL: SEED_DIR '$SEED_DIR' is not a directory" >&2; exit 1; }
  cp -a "$SEED_DIR"/. "$UDATA/warden/"
fi

mkdir -p "$SCRATCH/empty"

# mkfs an ext4 partition image of exactly $2 bytes from staged dir $1.
mkfs_part() {
  local stage="$1" bytes="$2" img="$3"
  rm -f "$img"
  truncate -s "$bytes" "$img"
  mkfs.ext4 -F -q -d "$stage" "$img"
}

DISK="$OUT/disk.img"
rm -f "$DISK"

# dd a REAL raw image into a partition window, fail-closed on overflow.
place_real_image() { # $1 src image, $2 offset, $3 partition size, $4 name
  local srcsz
  srcsz=$(stat -c %s "$1")
  [ "$srcsz" -le "$3" ] || {
    echo "FATAL: $4 image $1 ($srcsz bytes) exceeds the $3-byte partition" >&2
    exit 1
  }
  dd if="$1" of="$DISK" bs=4096 seek=$(($2 / 4096)) conv=notrunc,sparse status=none
  qemu_log "  $4: REAL image $(basename "$1") ($((srcsz / 1048576))M) @ $2"
}

place_partition() {
  local name="$1" off="$2" size="$3" stage=""
  case "$name" in
    rootfs_a)
      if [ -n "$ROOTFS_IMAGE" ]; then
        [ -f "$ROOTFS_IMAGE" ] || { echo "FATAL: --rootfs-image $ROOTFS_IMAGE not found" >&2; exit 1; }
        DISK_END_TRACK "$off" "$size"
        place_real_image "$ROOTFS_IMAGE" "$off" "$size" "$name"
        return 0
      fi
      stage="$ROOT" ;;
    oem_a)
      if [ -n "$OEM_IMAGE" ]; then
        [ -f "$OEM_IMAGE" ] || { echo "FATAL: --oem-image $OEM_IMAGE not found" >&2; exit 1; }
        DISK_END_TRACK "$off" "$size"
        place_real_image "$OEM_IMAGE" "$off" "$size" "$name"
        return 0
      fi
      stage="$SCRATCH/empty" ;;
    rootfs_b) stage="$ROOT" ;;
    userdata)          stage="$UDATA" ;;
    oem_b)             stage="$SCRATCH/empty" ;;
    # misc carries REAL AvbABData (byte 2048): flared's slotctl fail-closes
    # on a bad magic before any OTA write, so a zeroed misc blocks apply
    # scenarios. Bytes mirror flare-edge tools/mk-misc.py provisioning
    # defaults (A: prio 15 successful, B: prio 14 successful, CRC32-BE).
    misc) stage="__misc__" ;;
    # Boot-chain partitions the VM never reads: present at the right offsets,
    # left zeroed. Enumerated (not a wildcard) so a typo'd name in
    # blkdevparts.conf fails HERE, not as a confusing mount error at boot.
    env|idblock|uboot|boot_a|boot_b|recovery) stage="" ;;
    *) echo "FATAL: unknown partition name '$name' in blkdevparts.conf" >&2; exit 1 ;;
  esac
  DISK_END_TRACK "$off" "$size"
  # dd in 4K blocks: every offset in the canonical layout is 4K-aligned;
  # assert rather than assume, a misaligned write would corrupt a neighbor.
  if [ $((off % 4096)) -ne 0 ] || [ $((size % 4096)) -ne 0 ]; then
    echo "FATAL: partition $name not 4K-aligned (off=$off size=$size)" >&2
    exit 1
  fi
  [ -z "$stage" ] && return 0
  local img="$SCRATCH/$name.img"
  if [ "$stage" = "__misc__" ]; then
    python3 - "$img" "$size" <<'PYMISC'
import struct, sys, zlib
img, size = sys.argv[1], int(sys.argv[2])
s = bytearray(28)
s[0:4] = b"\0AB0"          # AB_MAGIC
s[4] = 1                     # major
s[8:12] = bytes([15, 0, 1, 0])   # slot A: priority, tries, successful
s[12:16] = bytes([14, 0, 1, 0])  # slot B
meta = bytes(s) + struct.pack(">I", zlib.crc32(bytes(s)) & 0xFFFFFFFF)
buf = bytearray(size)
buf[2048:2048 + len(meta)] = meta
open(img, "wb").write(buf)
PYMISC
    dd if="$img" of="$DISK" bs=4096 seek=$((off / 4096)) \
       conv=notrunc,sparse status=none
    qemu_log "  $name: AvbABData provisioned @ +2048"
    return 0
  fi
  mkfs_part "$stage" "$size" "$img"
  dd if="$img" of="$DISK" bs=4096 seek=$((off / 4096)) \
     conv=notrunc,sparse status=none
  qemu_log "  $name: ext4, $((size / 1048576))M @ $off"
}

DISK_END=0
# Max, not last: blkdevparts grammar permits explicit @offsets out of order.
DISK_END_TRACK() { [ $(($1 + $2)) -gt "$DISK_END" ] && DISK_END=$(($1 + $2)); return 0; }
qemu_log "building $DISK ($WARDEN_BLKDEVPARTS)"
truncate -s 0 "$DISK"
qemu_each_partition place_partition
truncate -s "$DISK_END" "$DISK"
qemu_log "disk image: $DISK ($(du -h "$DISK" | cut -f1) used, $((DISK_END / 1048576))M apparent)"