#!/usr/bin/env bash
# Install a LexiSuite image to a USB stick (Linux).
#
#   sudo bash install-to-usb.sh                                   # newest LexiSuite-*.img.xz next to this script or here
#   sudo bash install-to-usb.sh LexiSuite-c4ce6d8-secureboot-dev.img.xz
#   sudo bash install-to-usb.sh https://.../LexiSuite-c4ce6d8-secureboot-dev.img.xz
#   sudo bash install-to-usb.sh IMAGE /dev/sdX                    # name the stick yourself
#   bash install-to-usb.sh --dry-run [IMAGE] [/dev/sdX]           # every check, nothing written, no root needed
#
# What it does: checks the image against its .sha256 file, offers only USB disks, refuses a disk that holds a mounted
# system folder or is too small, asks you to type the disk's name, unmounts it, writes the image, clears the old
# backup partition table at the end of the stick, reads the whole image back from the stick and compares it, then
# ejects. Everything on the stick is erased.
# Windows or macOS: write the .img.xz with balenaEtcher or Raspberry Pi Imager instead.
set -euo pipefail
umask 022

DRY=0
ARGS=()
for arg in "$@"; do
    case "$arg" in
        --dry-run|--check) DRY=1 ;;
        -h|--help) sed -n '2,15p' "$0"; exit 0 ;;
        *) ARGS+=("$arg") ;;
    esac
done
SRC=${ARGS[0]:-}
DEV=${ARGS[1]:-}
# Test hook: write into this regular file instead of a stick (skips the USB disk checks only).
TEST_TARGET=${LEXI_TEST_TARGET:-}

die() { echo "Stopped: $*" >&2; echo 'Nothing was written.' >&2; exit 1; }
need() { command -v "$1" >/dev/null 2>&1 || die "this needs the '$1' command (install it and run again)"; }
for tool in lsblk dd cmp sha256sum xz awk; do need "$tool"; done

# ---------------------------------------------------------------- the image
WORK=""
cleanup() { [ -n "$WORK" ] && rm -rf -- "$WORK"; }
trap cleanup EXIT
if [ -z "$SRC" ]; then
    here=$(cd "$(dirname "$0")" && pwd)
    SRC=$(ls -t "$here"/LexiSuite-*.img.xz ./LexiSuite-*.img.xz "$here"/LexiSuite-*.img ./LexiSuite-*.img 2>/dev/null | head -n 1 || true)
    [ -n "$SRC" ] || die "no LexiSuite-*.img.xz found next to this script or in $(pwd); give the image file or its link"
fi
case "$SRC" in
    http://*|https://*)
        need curl
        WORK=$(mktemp -d)
        IMG="$WORK/$(basename "${SRC%%\?*}")"
        echo "Downloading $SRC"
        curl -fL --retry 3 -o "$IMG" "$SRC" || die "download failed"
        curl -fsL --retry 3 -o "$IMG.sha256" "$SRC.sha256" || die "the checksum file $SRC.sha256 could not be downloaded"
        ;;
    *) IMG=$SRC ;;
esac
[ -f "$IMG" ] || die "image not found: $IMG"
[ -f "$IMG.sha256" ] || die "checksum file $IMG.sha256 not found (it is published next to the image)"
echo "Checking $(basename "$IMG") against its checksum..."
want=$(awk '{print $1; exit}' "$IMG.sha256")
have=$(sha256sum "$IMG" | awk '{print $1}')
[ "$want" = "$have" ] || die "the image does not match $IMG.sha256 (damaged or incomplete download)"
echo "  checksum OK"
case "$IMG" in
    *.xz) SIZE=$(xz --robot --list "$IMG" | awk '$1 == "totals" {print $5}')
          xz -t "$IMG" || die "the image archive is damaged" ;;
    *)    SIZE=$(stat -c %s "$IMG") ;;
esac
[ -n "$SIZE" ] && [ "$SIZE" -gt 0 ] || die "could not read the image size"
stream() { case "$IMG" in *.xz) xz -dc "$IMG" ;; *) cat "$IMG" ;; esac; }
echo "  image size $((SIZE / 1048576)) MiB"

# ---------------------------------------------------------------- the stick
usb_disks() { lsblk -d -n -b -p -o NAME,TRAN,TYPE,SIZE,MODEL 2>/dev/null | awk '$2 == "usb" && $3 == "disk" && $4 > 0'; }
# A stick a desktop mounted for you sits under /media or /run/media. A USB disk mounted anywhere else is a data
# drive in use (for example /mnt/backup) and is never offered or written.
foreign_mounts() { lsblk -n -o MOUNTPOINT "$1" | awk 'NF && $0 !~ "^(/media/|/run/media/)" {print}' | paste -sd ' ' -; }
if [ -n "$TEST_TARGET" ]; then
    DEV=$TEST_TARGET
    [ -f "$DEV" ] || die "test target $DEV is not a regular file"
    STICK=$(stat -c %s "$DEV"); MODEL="test file"
else
    if [ -z "$DEV" ]; then
        list=""
        echo; echo "USB disks on this computer:"
        while read -r name _ _ bytes model; do
            [ -n "$name" ] || continue
            busy=$(foreign_mounts "$name")
            if [ -n "$busy" ]; then
                printf '  %-12s %7.1f GB %-24s skipped: in use at %s\n' "$name" "$(awk -v b="$bytes" 'BEGIN{print b/1e9}')" "$model" "$busy"
            else
                printf '  %-12s %7.1f GB %s\n' "$name" "$(awk -v b="$bytes" 'BEGIN{print b/1e9}')" "$model"
                list="$list$name"$'\n'
            fi
        done <<< "$(usb_disks)"
        list=$(printf '%s' "$list" | sed '/^$/d')
        [ -n "$list" ] || die "no USB stick that is free to erase was found; plug the stick in and run again"
        if [ "$(echo "$list" | wc -l)" -eq 1 ]; then
            DEV=$list
        else
            read -r -p "Which one is the stick to erase (for example /dev/sdb)? " DEV || die "no stick chosen"
        fi
    fi
    case "$DEV" in /dev/*) ;; *) DEV=/dev/$DEV ;; esac
    DEV=$(readlink -f "$DEV")
    [ -b "$DEV" ] || die "$DEV is not a disk"
    [ "$(lsblk -d -n -o TYPE "$DEV")" = disk ] || die "$DEV is a partition; give the whole disk (for example /dev/sdb, not /dev/sdb1)"
    [ "$(lsblk -d -n -o TRAN "$DEV")" = usb ] || die "$DEV is not a USB disk; this script only writes USB sticks"
    busy=$(foreign_mounts "$DEV")
    [ -z "$busy" ] || die "$DEV is in use at $busy; that is a data drive, not a stick to erase"
    STICK=$(lsblk -d -n -b -o SIZE "$DEV")
    MODEL=$(lsblk -d -n -o VENDOR,MODEL "$DEV" | tr -s ' ')
fi
[ "$STICK" -ge "$SIZE" ] || die "the stick ($((STICK / 1048576)) MiB) is smaller than the image ($((SIZE / 1048576)) MiB)"
echo
echo "Stick: $DEV ($MODEL, $((STICK / 1000000000)) GB)"
echo "Everything on it will be erased and replaced by $(basename "$IMG")."
if [ "$DRY" = 1 ]; then echo "Dry run: every check passed; stopping before anything is written."; exit 0; fi
[ -n "$TEST_TARGET" ] || [ "$(id -u)" = 0 ] || die "run it with sudo to write the stick"
short=$(basename "$DEV")
read -r -p "Type $short to erase it and continue: " answer || die "no answer"
[ "$answer" = "$short" ] || die "you typed '$answer', not '$short'"

# ---------------------------------------------------------------- write, read back, eject
if [ -z "$TEST_TARGET" ]; then
    for part in $(lsblk -n -p -o NAME,MOUNTPOINT "$DEV" | awk 'NF == 2 {print $1}'); do
        echo "Unmounting $part"; umount "$part" || die "could not unmount $part (close any window using it)"
    done
    DIRECT=oflag=direct
else
    DIRECT=
fi
echo "Writing... (a few minutes on most sticks)"
stream | dd of="$DEV" bs=4M iflag=fullblock $DIRECT conv=fsync,notrunc status=progress
# A larger stick still has its old backup partition table in its last sectors; clear it.
if [ "$STICK" -gt "$SIZE" ] && [ "$STICK" -ge 1048576 ]; then
    dd if=/dev/zero of="$DEV" bs=1048576 count=1 seek=$((STICK / 1048576 - 1)) conv=fsync,notrunc status=none
fi
sync
[ -n "$TEST_TARGET" ] || blockdev --flushbufs "$DEV" 2>/dev/null || true
echo "Reading it back..."
stream | cmp -n "$SIZE" - "$DEV" || { echo "Stopped: the stick does not read back what was written. Try another stick or port." >&2; exit 1; }
echo "  read-back matches"
if [ -z "$TEST_TARGET" ]; then
    if command -v eject >/dev/null 2>&1; then eject "$DEV" 2>/dev/null || true
    elif command -v udisksctl >/dev/null 2>&1; then udisksctl power-off -b "$DEV" 2>/dev/null || true; fi
fi
cat <<'EOF'

Done. The stick is ready and safe to remove.

Booting it:
  1. Plug it into the PC, open the firmware boot menu (often F11, F12 or Esc at power-on) and pick the stick's
     UEFI entry. The LexiSuite menu offers 1 LexiKernel and 2 LexiOS on Linux.
  2. Secure Boot: this is a development image. Either add its certificate to the firmware's allowed list (db)
     once: Append the file EFI/LEXIOS/OWNER.DER from the stick's LEXIOS partition under Secure Boot key
     management; or boot with Secure Boot off and press Y at the development notice.
     Full steps: https://github.com/W61k3r/LexiSuite-Releases/wiki/Secure-Boot
  3. On first boot the data partition grows to fill the stick. The screen then shows the Panel address: open it
     from another computer on the same network and create the admin account.
EOF