#!/bin/sh
# validate-image.sh — functional test of an image produced by enable-ssh.sh.
#
# Extracts the output image's rootfs (p2) into a scratch tree (debugfs
# rdump — no loop device needed) inside a privileged Docker container,
# starts dropbear via the image's OWN /etc/init.d/sshd (which generates
# host keys into /data/dropbear exactly as on first boot), then:
#   1. checks the deployed authorized_keys equals the -k public key you gave
#      enable-ssh.sh (static check, outside the container)
#   2. ssh-login as root with a throwaway ed25519 key generated INSIDE the
#      container (no private key of yours is ever read or stored)
#   3. confirms password auth is refused
#   4. runs sftp put/get to prove the sftp-server wiring
#
# The image file itself is only read; all test mutations happen on a scratch
# copy of p2 inside the container.
#
# Usage: validate-image.sh -i OUTPUT_IMAGE -k PUBLIC_KEY [--image docker-image]
#
# Requires: docker (linux containers), a volume with ~1 GiB free for the p2
# scratch copy. Takes a few minutes (container setup + key generation).
set -eu

SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
INPUT= KEYFILE= DOCKER_IMAGE="debian:12"

while [ $# -gt 0 ]; do
  case "$1" in
    -i) INPUT=$2; shift 2 ;;
    -k) KEYFILE=$2; shift 2 ;;
    --image) DOCKER_IMAGE=$2; shift 2 ;;
    -h|--help) sed -n '2,30p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) echo "error: unknown argument: $1" >&2; exit 2 ;;
  esac
done

err() { echo "error: $*" >&2; exit 1; }
say() { echo "[validate] $*"; }

[ -n "$INPUT" ] && [ -n "$KEYFILE" ] || { err "usage: validate-image.sh -i OUTPUT_IMAGE -k PUBLIC_KEY"; }
[ -f "$INPUT" ] || err "image not found: $INPUT"
[ -f "$KEYFILE" ] || err "public key not found: $KEYFILE"
command -v docker >/dev/null 2>&1 || err "docker not found"
docker info >/dev/null 2>&1 || err "docker daemon not reachable"

# --- parse the MBR (read-only) for p2 ----------------------------------------
P2_INFO=$(python3 - "$INPUT" <<'PYEOF'
import struct, sys
with open(sys.argv[1], "rb") as f:
    mbr = f.read(512)
assert mbr[510:512] == b"\x55\xaa", "no MBR boot signature"
for i in range(4):
    e = mbr[446 + 16*i : 446 + 16*(i+1)]
    if e and e != b"\x00" * 16 and e[4] == 0x83:
        lba, size = struct.unpack("<II", e[8:16])
        if size:
            print(lba, size)
            break
else:
    sys.exit("no Linux (0x83) partition")
PYEOF
) || err "MBR parse failed"
P2_SECTORS=$(echo "$P2_INFO" | cut -d' ' -f1)
P2_SIZE_SECTORS=$(echo "$P2_INFO" | cut -d' ' -f2)
P2_OFFSET=$((P2_SECTORS * 512))
P2_SIZE=$((P2_SIZE_SECTORS * 512))

WORK=$(mktemp -d "${TMPDIR:-/tmp}/ssh-validate.XXXXXX")
trap 'rm -rf "$WORK"' EXIT
# keep the scratch under $HOME on macOS: Docker Desktop bind-mounts under /tmp
# are unreliable (see pinball-assets/README.md "Gotchas learned"). If $WORK is
# under /tmp, copy it to $HOME.
case "$WORK" in
  /tmp/*|/var/folders/*) WORK_HOME="$HOME/.ssh-validate.$$"; mkdir -p "$WORK_HOME"; trap 'rm -rf "$WORK" "$WORK_HOME"' EXIT; WORK="$WORK_HOME" ;;
esac

say "extracting p2 from $INPUT (read-only) -> $WORK/p2.img"
dd if="$INPUT" of="$WORK/p2.img" bs=512 skip="$P2_SECTORS" count="$P2_SIZE_SECTORS" 2>/dev/null
[ "$(wc -c < "$WORK/p2.img" | tr -d '[:space:]')" = "$P2_SIZE" ] || err "p2 extraction size mismatch"

# --- static check: deployed authorized_keys must equal the -k input -----------
say "static check: authorized_keys in the image vs -k file"
debugfs -R "dump /home/root/.ssh/authorized_keys $WORK/deployed-ak" "$WORK/p2.img" 2>/dev/null \
  || err "image has no /home/root/.ssh/authorized_keys (was it built with enable-ssh.sh?)"
# normalise: ignore trailing-newline differences
tr -d '\n' < "$WORK/deployed-ak" > "$WORK/deployed-ak.n"
tr -d '\n' < "$KEYFILE" > "$WORK/keyfile.n"
if ! cmp -s "$WORK/deployed-ak.n" "$WORK/keyfile.n"; then
  err "authorized_keys in the image does not match -k ($KEYFILE)"
fi
say "  authorized_keys matches"

# --- container test ------------------------------------------------------------
say "running container test ($DOCKER_IMAGE, ~a few minutes) ..."
cat > "$WORK/container-test.sh" <<'TEST'
#!/bin/sh
set -e
apt-get update -qq >/dev/null
apt-get install -y -qq openssh-client qemu-user-static e2fsprogs >/dev/null 2>&1
# extract the rootfs to a scratch tree with debugfs (no loop device needed —
# Docker Desktop for Mac does not provide one; rdump preserves symlinks and modes)
mkdir -p /mnt2
debugfs -R 'rdump / /mnt2' /work/p2.img 2>/dev/null
echo "--- rootfs extracted (scratch copy); sanity ---"
ls -la /mnt2/bin/sh /mnt2/etc/init.d/sshd /mnt2/usr/local/bin/dropbear
# chrooted ARM binaries need an explicit qemu prefix (no binfmt in the container)
cp /usr/bin/qemu-arm-static /mnt2/usr/bin/
mount --bind /dev /mnt2/dev
mount -t proc proc /mnt2/proc
# start dropbear exactly the way the machine will, via the image's own init
# script (first-boot path: generates host keys into /data/dropbear, then runs
# the daemon on port 22)
echo "--- /etc/init.d/sshd start ---"
timeout 300 chroot /mnt2 /usr/bin/qemu-arm-static /bin/sh /etc/init.d/sshd start
sleep 3
ls -la /mnt2/data/dropbear/
echo "--- pubkey login test (throwaway key, user root) ---"
ssh-keygen -t ed25519 -N '' -f /tmp/vkey -q
mkdir -p /mnt2/home/root/.ssh
cp /tmp/vkey.pub /mnt2/home/root/.ssh/authorized_keys
chmod 700 /mnt2/home/root/.ssh; chmod 600 /mnt2/home/root/.ssh/authorized_keys
if ! timeout 90 ssh -i /tmp/vkey -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    -p 22 root@127.0.0.1 'echo LOGIN-OK as $(whoami) on $(hostname); uname -m' > /tmp/sshout 2>&1; then
  echo "FAIL: pubkey login did not succeed"; cat /tmp/sshout; exit 1
fi
grep -v Warning /tmp/sshout
echo "--- password auth must be refused ---"
if timeout 30 ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
     -o PubkeyAuthentication=no -o NumberOfPasswordPrompts=0 \
     -p 22 root@127.0.0.1 'echo SHOULD-NOT-PRINT' 2>/dev/null; then
  echo "FAIL: password auth was accepted"; exit 1
fi
echo "  (refused, as expected)"
echo "--- sftp put/get test ---"
echo hello-spike2 > /work/sftp-test.txt
if ! timeout 60 sftp -P 22 -i /tmp/vkey -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
     root@127.0.0.1 <<'SFTP' >/dev/null 2>&1
put /work/sftp-test.txt /tmp/sftp-test.txt
get /tmp/sftp-test.txt /work/sftp-back.txt
SFTP
then
  echo "FAIL: sftp session failed"
  exit 1
fi
cmp /work/sftp-test.txt /work/sftp-back.txt
echo "  sftp round-trip OK"
echo "=== VALIDATE-ALL-OK ==="
exit 0
TEST
chmod +x "$WORK/container-test.sh"

if ! docker run --rm --privileged \
  -v "$WORK:/work:rw" \
  "$DOCKER_IMAGE" \
  /bin/sh /work/container-test.sh; then
  err "container test failed"
fi

say "PASS: key-only sshd boots from the image's own init script, pubkey login works, password auth refused, sftp works."
say "(scratch p2 and container artifacts deleted; your image file was only read)"