#!/bin/sh
# Write schos.iso to a USB stick.
#
#   ./write.sh sdc            or   ./write.sh /dev/sdc
#   ./write.sh sdc other.iso  to write a different image
#   ./write.sh sdc --no-verify  skip the read-back check (halves the time)
#
# Shows what is on the device, asks you to type the device's name back, then
# asks for your password and writes. Everything on the device is destroyed;
# there is no undo.
#
# Linux only: it reads the device list from lsblk. It refuses a partition
# (sdc1 -- the stick is written whole), anything holding a mounted system
# path, and anything that is not a removable or USB disk unless you add
# --any-disk, because the one mistake this script exists to stop is writing
# over the disk you are running from.

set -eu

if [ -t 1 ]; then
    R=$(printf '\033[1;31m'); G=$(printf '\033[1;32m')
    B=$(printf '\033[1m'); D=$(printf '\033[2m'); N=$(printf '\033[0m')
else
    R='' G='' B='' D='' N=''
fi

die() { printf '%swrite.sh:%s %s\n' "$R" "$N" "$*" >&2; exit 1; }

usage() {
    echo "usage: ./write.sh <device> [image.iso] [--any-disk] [--no-verify]"
    echo "       ./write.sh sdc        writes schos.iso to /dev/sdc"
    echo
    echo "Removable disks on this machine:"
    lsblk -d -n -o NAME,SIZE,TRAN,RM,MODEL 2>/dev/null |
        awk '$3 == "usb" || $4 == "1" { printf "  /dev/%-8s %8s  %s\n", $1, $2, substr($0, index($0,$5)) }'
    exit 2
}

ANY=0
VERIFY=1
DEV=
ISO=
for a in "$@"; do
    case "$a" in
        --any-disk) ANY=1 ;;
        --no-verify) VERIFY=0 ;;
        -h|--help)  usage ;;
        *) if [ -z "$DEV" ]; then DEV=$a; elif [ -z "$ISO" ]; then ISO=$a; else usage; fi ;;
    esac
done
[ -n "$DEV" ] || usage

HERE=$(cd "$(dirname "$0")" && pwd)
[ -n "$ISO" ] || ISO=$HERE/schos.iso

command -v lsblk >/dev/null 2>&1 || die "needs lsblk (util-linux); this script is Linux only"

# --- the image --------------------------------------------------------------

[ -f "$ISO" ] || die "no image at $ISO -- run 'make iso' first"
ISO_BYTES=$(wc -c < "$ISO" | tr -d ' ')
ISO_DATE=$(date -r "$ISO" '+%Y-%m-%d %H:%M' 2>/dev/null || echo '?')

# --- the device -------------------------------------------------------------

case "$DEV" in
    /dev/*) ;;
    *)      DEV=/dev/$DEV ;;
esac
[ -b "$DEV" ] || die "$DEV is not a block device"

NAME=$(basename "$DEV")
TYPE=$(lsblk -d -n -o TYPE "$DEV")
[ "$TYPE" = disk ] || die "$DEV is a $TYPE, not a whole disk -- give the disk (sdc), not a partition (sdc1)"

SIZE_BYTES=$(lsblk -d -n -b -o SIZE "$DEV" | tr -d ' ')
SIZE=$(lsblk -d -n -o SIZE "$DEV" | tr -d ' ')
TRAN=$(lsblk -d -n -o TRAN "$DEV" | tr -d ' ')
RM=$(lsblk -d -n -o RM "$DEV" | tr -d ' ')
MODEL=$(lsblk -d -n -o MODEL "$DEV" | sed 's/ *$//')

# The disk this system runs from, or anything with a system path on it.
for mp in $(lsblk -n -o MOUNTPOINTS "$DEV" 2>/dev/null); do
    case "$mp" in
        /|/boot|/boot/*|/efi|/home|/nix|/nix/*|/usr|/var|\[SWAP\])
            die "$DEV holds $mp -- that is this computer's own disk, not a stick" ;;
    esac
done

if [ "$TRAN" != usb ] && [ "$RM" != 1 ] && [ "$ANY" != 1 ]; then
    die "$DEV ($MODEL, $SIZE, ${TRAN:-no transport}) is not a USB or removable disk.
          If you are sure, run it again with --any-disk."
fi

[ "$SIZE_BYTES" -ge "$ISO_BYTES" ] || die "$DEV is $SIZE, smaller than the image"

# --- show, and ask ----------------------------------------------------------

echo
echo "${B}Image${N}   $ISO"
echo "        $((ISO_BYTES / 1048576)) MB, built $ISO_DATE"
echo
echo "${B}Device${N}  $DEV  ${B}$MODEL${N}  $SIZE  ${D}(${TRAN:-?}, removable=$RM)${N}"
lsblk -n -o NAME,SIZE,FSTYPE,LABEL,MOUNTPOINTS "$DEV" | sed '1d; s/^/        /'
USED=$(lsblk -n -o NAME "$DEV" | sed 1d | wc -l | tr -d ' ')
[ "$USED" -gt 0 ] || echo "        ${D}(no partitions)${N}"
echo
echo "${R}Everything on $DEV will be destroyed. This cannot be undone.${N}"
printf 'Type %s%s%s to write, anything else to stop: ' "$B" "$NAME" "$N"
read -r answer
[ "$answer" = "$NAME" ] || { echo "Stopped. Nothing was written."; exit 1; }

# --- write ------------------------------------------------------------------

echo
echo "Your password is needed to write to the device."
sudo -v || die "no permission; nothing was written"

# Unmount whatever the desktop mounted, or the kernel keeps writing its own
# cached view of the old filesystem over the new image afterwards.
for part in $(lsblk -n -l -o NAME,MOUNTPOINTS "$DEV" | awk 'NF > 1 { print $1 }'); do
    echo "Unmounting /dev/$part"
    sudo umount "/dev/$part" || die "could not unmount /dev/$part; nothing was written"
done

echo "Writing..."
sudo dd if="$ISO" of="$DEV" bs=4M conv=fsync oflag=direct status=progress
sync

# --- check ------------------------------------------------------------------

# Read it back. A stick that takes the write and keeps none of it is common
# enough among cheap ones to be worth the time.
#
# iflag=direct and 4 MB reads: a plain `cmp ISO /dev/sdX` reads through the page
# cache in small requests, which a USB stick serves slowly -- and may be answered
# from the cache rather than from the stick, checking nothing. Direct reads
# check the stick itself, and at the speed it can actually deliver.
if [ "$VERIFY" = 1 ]; then
    echo "Checking what was written..."
    BLOCKS=$(( (ISO_BYTES + 4194303) / 4194304 ))
    if sudo dd if="$DEV" bs=4M count="$BLOCKS" iflag=direct status=none | cmp -n "$ISO_BYTES" - "$ISO"; then
        echo "${G}Done.${N} $DEV holds $(basename "$ISO") and can be removed."
    else
        die "the device does not read back what was written -- the stick may be failing"
    fi
else
    echo "${G}Done.${N} $DEV holds $(basename "$ISO") (not verified) and can be removed."
fi

# The kernel re-reads the partition table so the new layout shows up here
# too; harmless if it cannot.
sudo partprobe "$DEV" 2>/dev/null || sudo blockdev --rereadpt "$DEV" 2>/dev/null || true