#!/usr/bin/env bash
# write_stick.sh -- write a StarshipOS stick image to a USB stick, safely, and prove it.
#
#   scripts/write_stick.sh IMAGE DEVICE        e.g. build/boards/pc/lithos-pc.img /dev/sdb
#   make write_stick DEVICE=/dev/sdX           (the PC stick; root Makefile)
#
# Linux only (lsblk, udisksctl, sudo). Refuses a device that is not a whole, removable USB disk,
# that holds the running system, that is smaller than the image, or that has a mounted partition.
# Shows the device and asks for its name to be typed back before writing. Writes with
# conv=fsync, flushes, reads the stick back and compares it with the image byte for byte, then
# powers the stick off so it can be pulled. Everything on the stick is erased.
#
# On Windows or macOS: balenaEtcher (https://etcher.balena.io) writes and verifies an image;
# on Windows, Rufus in "DD Image" mode does too.
set -euo pipefail
PATH="$PATH:/usr/sbin:/sbin"

die() { echo "write_stick: $*" >&2; exit 1; }

[ $# -eq 2 ] || die "usage: scripts/write_stick.sh IMAGE DEVICE   (e.g. /dev/sdb; see lsblk)"
IMG="$1" DEV="$2"
[ -f "$IMG" ] || die "no image at $IMG -- build it first: make boot_image TARGET=PC"
[ -b "$DEV" ] || die "$DEV is not a block device (lsblk lists the disks)"

name=$(basename "$DEV")
[ "$(lsblk -dno TYPE "$DEV")" = disk ] || die "$DEV is a partition, not a whole disk: give the disk (e.g. /dev/sdb, not /dev/sdb1)"
[ "$(lsblk -dno RM "$DEV" | tr -d ' ')" = 1 ] || die "$DEV is not a removable disk -- refusing"
[ "$(lsblk -dno TRAN "$DEV" | tr -d ' ')" = usb ] || die "$DEV is not a USB disk -- refusing"
root_src=$(findmnt -no SOURCE / || true)
if [ -n "$root_src" ] && lsblk -no PKNAME "$root_src" 2>/dev/null | grep -qx "$name"; then
    die "$DEV holds the running system -- refusing"
fi
img_bytes=$(stat -c %s "$IMG")
dev_bytes=$(lsblk -bdno SIZE "$DEV")
[ "$dev_bytes" -ge "$img_bytes" ] || die "$DEV ($dev_bytes bytes) is smaller than the image ($img_bytes bytes)"
mounted=$(lsblk -nro NAME,MOUNTPOINTS "$DEV" | awk 'NF > 1 {print "/dev/" $1}')
if [ -n "$mounted" ]; then
    echo "write_stick: $DEV has mounted partitions. Unmount them first:" >&2
    for p in $mounted; do echo "  udisksctl unmount -b $p" >&2; done
    exit 1
fi

echo "About to ERASE this whole stick and write the StarshipOS image to it:"
lsblk -o NAME,SIZE,TRAN,MODEL,VENDOR,FSTYPE,LABEL "$DEV"
echo
echo "  image: $IMG ($img_bytes bytes)"
read -r -p "Type the device name ($DEV) to go on, anything else to stop: " answer
[ "$answer" = "$DEV" ] || die "stopped; nothing was written"

echo "== writing (sudo asks for your password)"
sudo dd if="$IMG" of="$DEV" bs=4M conv=fsync status=progress
echo "== flushing to the stick"
sudo sync
sudo blockdev --flushbufs "$DEV"
echo "== reading the stick back and comparing it with the image"
if sudo cmp -n "$img_bytes" "$IMG" "$DEV"; then
    echo "== the stick matches the image"
else
    die "the stick does NOT match the image -- do not boot it; try again, or another stick"
fi
if command -v udisksctl >/dev/null 2>&1 && udisksctl power-off -b "$DEV" >/dev/null 2>&1; then
    echo "== done: the stick is powered off and can be pulled"
else
    echo "== done: the stick can be pulled"
fi
echo "Boot it: Secure Boot off in the firmware setup, then choose the stick in the boot menu."