#!/bin/bash
set -Eeuo pipefail

if [ "$(id -u)" -ne 0 ]; then
    printf '%s\n' 'verify_image_bootstrap.sh must run as root.' >&2
    exit 1
fi
if [ "$#" -ne 5 ]; then
    printf '%s\n' 'usage: verify_image_bootstrap.sh IMAGE SERVICE_UNIT SSH_CONFIG KERNEL_DEPENDENCY APP_BUNDLE' >&2
    exit 2
fi

IMAGE=$(readlink -f -- "$1")
UNIT=$(readlink -f -- "$2")
SSH_CONFIG=$(readlink -f -- "$3")
KERNEL_DEPENDENCY=$(readlink -f -- "$4")
APP_BUNDLE=$(readlink -f -- "$5")
SCRIPT_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
[ -f "$IMAGE" ]
[ -f "$UNIT" ]
[ -f "$SSH_CONFIG" ]
[ -d "$KERNEL_DEPENDENCY" ]
[ -f "$APP_BUNDLE" ]
[ -f "$KERNEL_DEPENDENCY/SHA256SUMS" ]
[ -f "$KERNEL_DEPENDENCY/METADATA.json" ]

LOOP=
MOUNT_DIR=$(mktemp -d /var/tmp/matrix-image-verify.XXXXXX)
BOOT_MOUNT=$(mktemp -d /var/tmp/matrix-image-boot-verify.XXXXXX)
cleanup() {
    if mountpoint -q "$BOOT_MOUNT"; then
        umount "$BOOT_MOUNT"
    fi
    if mountpoint -q "$MOUNT_DIR"; then
        umount "$MOUNT_DIR"
    fi
    if [ -n "$LOOP" ]; then
        losetup -d "$LOOP"
    fi
    rmdir "$MOUNT_DIR" 2>/dev/null || true
    rmdir "$BOOT_MOUNT" 2>/dev/null || true
}
trap cleanup EXIT

LOOP=$(losetup --find --show --partscan --read-only "$IMAGE")
ROOT_PARTITION="${LOOP}p2"
BOOT_PARTITION="${LOOP}p1"
for _attempt in $(seq 1 20); do
    [ -b "$ROOT_PARTITION" ] && break
    sleep 0.25
done
[ -b "$ROOT_PARTITION" ]
[ -b "$BOOT_PARTITION" ]
mount -o ro,noload "$ROOT_PARTITION" "$MOUNT_DIR"
mount -o ro "$BOOT_PARTITION" "$BOOT_MOUNT"

INSTALLED="$MOUNT_DIR/etc/systemd/system/matrix-image-firstboot.service"
ENABLED="$MOUNT_DIR/etc/systemd/system/multi-user.target.wants/matrix-image-firstboot.service"
SSH_ENABLED="$MOUNT_DIR/etc/systemd/system/multi-user.target.wants/ssh.service"
SSH_INSTALLED="$MOUNT_DIR/etc/ssh/sshd_config.d/10-walnutpi-screen-hardening.conf"
cmp -s "$UNIT" "$INSTALLED"
cmp -s "$SSH_CONFIG" "$SSH_INSTALLED"
[ -L "$ENABLED" ]
[ "$(readlink "$ENABLED")" = '../matrix-image-firstboot.service' ]
grep -Fxq 'ExecStart=/boot/MSCINIT' "$INSTALLED"
grep -Eq '^Before=.*ssh\.service.*boot_script\.service' "$INSTALLED"
[ -x "$MOUNT_DIR/usr/sbin/sshd" ]
[ -x "$MOUNT_DIR/usr/bin/sudo" ]
[ -x "$MOUNT_DIR/usr/sbin/visudo" ]
[ -f "$MOUNT_DIR/lib/systemd/system/ssh.service" ]
[ -L "$SSH_ENABLED" ]
[ "$(basename "$(readlink "$SSH_ENABLED")")" = 'ssh.service' ]
[ ! -L "$MOUNT_DIR/etc/systemd/system/ssh.service" ] || [ "$(readlink "$MOUNT_DIR/etc/systemd/system/ssh.service")" != '/dev/null' ]
[ ! -s "$MOUNT_DIR/etc/machine-id" ]
[ "$(readlink "$MOUNT_DIR/var/lib/dbus/machine-id")" = '/etc/machine-id' ]
if find "$MOUNT_DIR/etc/ssh" -maxdepth 1 -name 'ssh_host_*' -print -quit 2>/dev/null | grep -q .; then
    printf '%s\n' 'export still contains an SSH host key' >&2
    exit 1
fi
if find "$MOUNT_DIR/etc/NetworkManager/system-connections" -mindepth 1 -maxdepth 1 -type f -print -quit 2>/dev/null | grep -q .; then
    printf '%s\n' 'export still contains a NetworkManager connection' >&2
    exit 1
fi

KERNEL_ARCHIVE=axp313a-kernel-6.1.31-matrix-axp313a1.tar.gz
KERNEL_TARGET=$MOUNT_DIR/opt/matrix-image-bootstrap/axp313a
APP_TARGET=$MOUNT_DIR/opt/matrix-image-bootstrap/app/MSCBOOT.TGZ
PYTHONPATH="$SCRIPT_ROOT" python3 "$SCRIPT_ROOT/scripts/image_bootstrap_payload.py" verify \
    --root "$MOUNT_DIR" --bundle "$APP_BUNDLE" --dependency "$KERNEL_DEPENDENCY" >/dev/null
[ ! -e "$BOOT_MOUNT/MSCBOOT.TGZ" ]
python3 - "$BOOT_MOUNT/MSCMETA.JSN" "$KERNEL_DEPENDENCY/METADATA.json" \
    "$KERNEL_TARGET/$KERNEL_ARCHIVE" "$APP_TARGET" <<'PY'
import hashlib
import json
from pathlib import Path
import sys

image = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
kernel = json.loads(Path(sys.argv[2]).read_text(encoding="utf-8"))
artifact = Path(sys.argv[3])
bundle = Path(sys.argv[4])
assert image["format_version"] == 3
assert image["bundle_path"] == "/opt/matrix-image-bootstrap/app/MSCBOOT.TGZ"
assert image["bundle_bytes"] == bundle.stat().st_size
assert image["kernel_release"] == kernel["kernel_release"]
assert image["kernel_artifact"] == kernel["artifact"]
assert image["kernel_artifact_bytes"] == kernel["artifact_bytes"] == artifact.stat().st_size
assert image["kernel_artifact_sha256"] == kernel["artifact_sha256"]
assert image["kernel_unpacked_file_bytes"] == kernel["unpacked_file_bytes"]
digest = hashlib.sha256()
with artifact.open("rb") as handle:
    for chunk in iter(lambda: handle.read(4 * 1024 * 1024), b""):
        digest.update(chunk)
assert digest.hexdigest() == kernel["artifact_sha256"]
digest = hashlib.sha256()
with bundle.open("rb") as handle:
    for chunk in iter(lambda: handle.read(4 * 1024 * 1024), b""):
        digest.update(chunk)
assert digest.hexdigest() == image["bundle_sha256"]
assert image["boot_safety_bytes"] == 32 * 1024 * 1024
assert image["boot_required_bytes"] > 0
PY

printf '%s\n' 'rootfs bootstrap verified'