#!/bin/sh
# enable-ssh.sh — inject key-only Dropbear SSH into a Stern Spike 2 raw SD image.
#
# Reproduces the working setup documented in pinball-assets/README.md (sections 1
# and 3): the ext3 rootfs partition (p2) of a factory Spike 2 image (King Kong
# LE 0.97.0, 16 GiB) is replaced with a copy that additionally contains
#   /usr/local/bin/dropbear         static musl armhf, 2025.88 — password auth
#                                   compiled out, DROPBEAR_REEXEC=0 (kernel 3.14
#                                   predates execveat)
#   /usr/local/bin/dropbearkey      static musl armhf
#   /usr/libexec/sftp-server       static Go sftp server — dropbear's compile-time
#                                   SFTPSERVER_PATH default, so scp/rsync work
#                                   without a -s flag
#   /etc/init.d/sshd               SysV init: generates host keys into
#                                   /data/dropbear at first boot, dropbear on 22
#   /etc/rc5.d/S30sshd             boot symlink -> ../init.d/sshd (before S95game)
#   /home/root/.ssh/authorized_keys  the public key(s) you supply (dir 0700, key 0600)
# Everything else in the image — p1 boot, the game partition, and all data/dump
# bytes — is left bit-for-bit untouched and verified to be so afterwards.
# No private keys are ever read, stored, or embedded.
#
# Usage:
#   enable-ssh.sh -i INPUT_IMAGE -o OUTPUT_IMAGE -k PUBLIC_KEY [options]
#
#   -i FILE   input raw SD image (factory/Release image). NEVER modified.
#   -o FILE   output raw SD image to create. Must differ from the input.
#   -k FILE   authorized_keys content: one or more public key lines
#             (ssh-ed25519, ssh-rsa, ecdsa-sha2-*, sk-*). No private keys.
#   -b DIR    dir with dropbear/dropbearkey/sftp-server (default: ./bin)
#   -f DIR    dir with the sshd init script (default: ./files)
#       --dry-run    parse and print the plan; write nothing
#       --force      allow a p2 geometry that differs from the known
#                    King Kong LE 0.97 layout (12 MiB offset, 336 MiB size)
#       --overwrite  allow an existing output file to be replaced
#   -h         help
#
# Exit status: 0 on success (or a complete dry run), non-zero on any error.
set -eu

SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
BINDIR="$SCRIPT_DIR/bin"
FILEDIR="$SCRIPT_DIR/files"

INPUT= OUTPUT= KEYFILE= DRYRUN=0 FORCE=0 OVERWRITE=0

usage() {
  cat <<'USAGE'
Usage: enable-ssh.sh -i INPUT_IMAGE -o OUTPUT_IMAGE -k PUBLIC_KEY [options]
  -i FILE   input raw SD image (factory/Release image). NEVER modified.
  -o FILE   output raw SD image to create. Must differ from the input.
  -k FILE   authorized_keys content: one or more public key lines
            (ssh-ed25519, ssh-rsa, ecdsa-sha2-*, sk-*). No private keys.
  -b DIR    dir with dropbear/dropbearkey/sftp-server (default: ./bin)
  -f DIR    dir with the sshd init script (default: ./files)
      --dry-run    parse and print the plan; write nothing
      --force      allow a p2 geometry different from the known layout
      --overwrite  allow an existing output file to be replaced
  -h         help
USAGE
}

while [ $# -gt 0 ]; do
  case "$1" in
    -i) INPUT=$2; shift 2 ;;
    -o) OUTPUT=$2; shift 2 ;;
    -k) KEYFILE=$2; shift 2 ;;
    -b) BINDIR=$2; shift 2 ;;
    -f) FILEDIR=$2; shift 2 ;;
    --dry-run) DRYRUN=1; shift ;;
    --force) FORCE=1; shift ;;
    --overwrite) OVERWRITE=1; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "error: unknown argument: $1" >&2; usage >&2; exit 2 ;;
  esac
done

err() { echo "error: $*" >&2; exit 1; }
say() { echo "[enable-ssh] $*"; }
sha256() {
  if command -v shasum >/dev/null 2>&1; then shasum -a 256 "$1" | cut -d' ' -f1
  else sha256sum "$1" | cut -d' ' -f1; fi
}
realpath() { python3 -c 'import os,sys; print(os.path.realpath(sys.argv[1]))' "$1"; }

[ -n "$INPUT" ]   || { usage >&2; err "missing -i INPUT_IMAGE"; }
[ -n "$OUTPUT" ]  || { usage >&2; err "missing -o OUTPUT_IMAGE"; }
[ -n "$KEYFILE" ] || { usage >&2; err "missing -k PUBLIC_KEY"; }

command -v python3 >/dev/null 2>&1 || err "python3 not found (needed for MBR parsing)"
command -v debugfs  >/dev/null 2>&1 || err "debugfs not found (brew install e2fsprogs)"
command -v e2fsck   >/dev/null 2>&1 || err "e2fsck not found (brew install e2fsprogs)"

# --- inputs -----------------------------------------------------------------
[ -f "$INPUT" ] && [ -r "$INPUT" ] || err "input image not a readable regular file: $INPUT"
IN_REAL=$(realpath "$INPUT")
OUT_REAL=$(realpath "$OUTPUT")
[ "$IN_REAL" != "$OUT_REAL" ] || err "output image must differ from input image ($OUT_REAL)"
if [ -e "$OUTPUT" ]; then
  [ "$OVERWRITE" = 1 ] || err "output exists: $OUTPUT (refusing to overwrite; use --overwrite)"
fi

for B in dropbear dropbearkey sftp-server; do
  [ -f "$BINDIR/$B" ] || err "missing binary: $BINDIR/$B (see README, 'Building the binaries')"
done
[ -f "$FILEDIR/sshd" ] || err "missing init script: $FILEDIR/sshd"

# public key: one or more public-key lines, nothing that smells private
[ -f "$KEYFILE" ] && [ -r "$KEYFILE" ] || err "public key file not readable: $KEYFILE"
KEYS=$(grep -c '[^[:space:]]' "$KEYFILE" 2>/dev/null) || KEYS=0
[ "$KEYS" -ge 1 ] || err "public key file has no lines: $KEYFILE"
grep -n 'PRIVATE' "$KEYFILE" >/dev/null 2>&1 && \
  err "key file contains 'PRIVATE' — refusing (private keys are not allowed here)"
if [ "$(tr -d '\r' < "$KEYFILE" | wc -c | tr -d '[:space:]')" != "$(wc -c < "$KEYFILE" | tr -d '[:space:]')" ]; then
  err "key file has CRLF line endings — convert with: tr -d '\\r' < in > out"
fi
while IFS= read -r line || [ -n "$line" ]; do
  case "$line" in
    '') continue ;;
    ssh-ed25519*|ssh-rsa*|ecdsa-sha2-nistp256*|ecdsa-sha2-nistp384*|ecdsa-sha2-nistp521*|sk-ssh-ed25519*|sk-ecdsa-sha2-nistp256*)
      ;;
    *) err "not a recognised public key line: $line" ;;
  esac
done < "$KEYFILE"

# --- partition table (read-only) ---------------------------------------------
# Known King Kong LE 0.97 layout: p1 FAT32 @4MiB, p2 ext3 rootfs @12MiB,
# 336 MiB (688128 sectors), then the game/data region. We parse the MBR
# rather than assume it.
P2_INFO=$(python3 - "$IN_REAL" <<'PYEOF'
import struct, sys
path = sys.argv[1]
with open(path, "rb") as f:
    mbr = f.read(512)
if mbr[510:512] != b"\x55\xaa":
    sys.exit("no MBR boot signature (55aa) — not a raw MBR image?")
for i in range(4):
    e = mbr[446 + 16*i : 446 + 16*(i+1)]
    if not e or e == b"\x00" * 16:
        continue
    if e[4] == 0x83:  # primary Linux partition
        lba, size = struct.unpack("<II", e[8:16])
        if size:
            print(lba, size)
            break
else:
    sys.exit("no Linux (0x83) partition in MBR")
PYEOF
) || err "MBR parse failed: $P2_INFO"
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))

KNOWN_OFFSET=12582912        # 12 MiB
KNOWN_SIZE=352321536         # 336 MiB
if [ "$P2_OFFSET" != "$KNOWN_OFFSET" ] || [ "$P2_SIZE" != "$KNOWN_SIZE" ]; then
  if [ "$FORCE" = 1 ]; then
    say "WARNING: p2 geometry differs from the known King Kong LE 0.97 layout"
    say "  parsed: offset $P2_OFFSET ($((P2_OFFSET / 1048576)) MiB), size $P2_SIZE ($((P2_SIZE / 1048576)) MiB)"
  else
    err "p2 geometry (offset $P2_OFFSET, size $P2_SIZE) differs from the known layout (offset $KNOWN_OFFSET, size $KNOWN_SIZE); re-run with --force to proceed"
  fi
fi

# input must be at least as large as the end of p2
IN_SIZE=$(wc -c < "$IN_REAL" | tr -d '[:space:]')
[ "$IN_SIZE" -ge $((P2_OFFSET + P2_SIZE)) ] || err "input image too small: $IN_SIZE bytes < p2 end $((P2_OFFSET + P2_SIZE))"

# p2 must be ext3 (magic 0xEF53, stored little-endian as 53 ef, at offset 1080)
MAGIC=$(python3 -c '
import sys
with open(sys.argv[1], "rb") as f:
    f.seek(int(sys.argv[2]))
    sys.stdout.write(f.read(2).hex())' "$IN_REAL" "$((P2_OFFSET + 1080))")
[ "$MAGIC" = "53ef" ] || err "p2 at offset $P2_OFFSET is not ext3 (magic $MAGIC, expected 53ef)"

say "input : $IN_REAL"
say "  size: $IN_SIZE bytes"
say "p2    : offset $P2_OFFSET ($((P2_OFFSET / 1048576)) MiB), size $P2_SIZE ($((P2_SIZE / 1048576)) MiB), ext3"
say "output: $OUT_REAL"
say "key   : $KEYFILE ($KEYS line(s))"

# --- dry run -----------------------------------------------------------------
if [ "$DRYRUN" = 1 ]; then
  say "dry run — nothing will be written"
  WORK=$(mktemp -d "${TMPDIR:-/tmp}/ssh-image.XXXXXX")
  trap 'rm -rf "$WORK"' EXIT
  say "extracting p2 from input for inspection (read-only) ..."
  dd if="$IN_REAL" of="$WORK/p2" bs=512 skip="$P2_SECTORS" count="$P2_SIZE_SECTORS" 2>/dev/null
  if debugfs -R 'stat /usr/local/bin/dropbear' "$WORK/p2" 2>/dev/null | grep -q '^Inode:'; then
    say "input p2 already contains /usr/local/bin/dropbear (image may already be SSH-enabled)"
  else
    say "input p2 does not contain dropbear (factory state)"
  fi
  say "plan:"
  say "  1. copy image: $IN_REAL -> $OUT_REAL (APFS clone if possible, else full copy)"
  say "  2. extract p2 from the copy, inject with debugfs (no mounting):"
  say "       /usr/local/bin/dropbear         <- $BINDIR/dropbear        (0755, 0:0)"
  say "       /usr/local/bin/dropbearkey      <- $BINDIR/dropbearkey     (0755, 0:0)"
  say "       /usr/libexec/sftp-server        <- $BINDIR/sftp-server     (0755, 0:0)"
  say "       /etc/init.d/sshd                <- $FILEDIR/sshd           (0755, 0:0)"
  say "       /etc/rc5.d/S30sshd              -> symlink ../init.d/sshd"
  say "       /home/root/.ssh                 (new dir, 0700, 0:0)"
  say "       /home/root/.ssh/authorized_keys <- $KEYFILE                (0600, 0:0)"
  say "  3. e2fsck -fn the injected p2 (abort if unhealthy)"
  say "  4. write p2 back at offset $P2_OFFSET"
  say "  5. verify: p2 region differs from input; every other byte identical (cmp)"
  exit 0
fi

# --- real run -----------------------------------------------------------------
WORK=$(mktemp -d "${TMPDIR:-/tmp}/ssh-image.XXXXXX")
trap 'rm -rf "$WORK"' EXIT
say "work dir: $WORK"

say "copying input image -> output (APFS clone if possible, else full copy) ..."
if ! cp -c "$IN_REAL" "$OUT_REAL" 2>/dev/null; then
  say "clone not available; doing a full copy (16 GiB — this takes a while)"
  cp "$IN_REAL" "$OUT_REAL"
fi
OUT_SIZE=$(wc -c < "$OUT_REAL" | tr -d '[:space:]')
[ "$OUT_SIZE" = "$IN_SIZE" ] || err "output size after copy mismatch: $OUT_SIZE != $IN_SIZE"

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

say "baseline filesystem check (e2fsck -fn) ..."
e2fsck -fn "$WORK/p2" >/dev/null 2>&1 || err "baseline e2fsck failed on p2 — input rootfs is not healthy; aborting (output left as an unmodified copy)"

# --- inject with debugfs (no mounting) ----------------------------------------
# debugfs `write` forces ownership 0:0 and copies the source file's mode; modes
# are pinned explicitly afterwards with `sif` (decimal 16-bit mode, type bits
# included: file 0755=33261, file 0600=33152, dir 0755=16877, dir 0700=16832).
# `symlink <name> <source>` requires the source to exist in the fs and resolves
# it against debugfs's cwd, so we `cd` into /etc/rc5.d first; the link content
# written is the literal source string.
say "injecting files with debugfs ..."
exists() { debugfs -R "stat $1" "$WORK/p2" 2>/dev/null | grep -q '^Inode:'; }

CMDS="$WORK/cmds"
: > "$CMDS"
exists /usr/libexec       || echo "mkdir /usr/libexec"        >> "$CMDS"
exists /home/root/.ssh    || echo "mkdir /home/root/.ssh"     >> "$CMDS"
exists /etc/rc5.d/S30sshd && echo "rm /etc/rc5.d/S30sshd"     >> "$CMDS"

echo "write $BINDIR/dropbear /usr/local/bin/dropbear"         >> "$CMDS"
echo "write $BINDIR/dropbearkey /usr/local/bin/dropbearkey"   >> "$CMDS"
echo "write $BINDIR/sftp-server /usr/libexec/sftp-server"     >> "$CMDS"
echo "write $FILEDIR/sshd /etc/init.d/sshd"                   >> "$CMDS"
echo "write $KEYFILE /home/root/.ssh/authorized_keys"         >> "$CMDS"
echo "cd /etc/rc5.d"                                          >> "$CMDS"
echo "symlink S30sshd ../init.d/sshd"                         >> "$CMDS"

DBG1="$WORK/dbg1.log"
debugfs -w -f "$CMDS" "$WORK/p2" > "$DBG1" 2>&1 || err "debugfs injection failed"
if grep -qiE 'error|failed|denied|not found' "$DBG1"; then
  sed 's/^/    /' "$DBG1" >&2
  err "debugfs reported errors (log above)"
fi

# pin the exact modes (session 2: non-interactive sif)
SMODES="$WORK/smodes"
: > "$SMODES"
echo "sif /usr/local/bin/dropbear         mode 33261" >> "$SMODES"  # 0100755
echo "sif /usr/local/bin/dropbearkey      mode 33261" >> "$SMODES"  # 0100755
echo "sif /usr/libexec/sftp-server        mode 33261" >> "$SMODES"  # 0100755
echo "sif /etc/init.d/sshd                mode 33261" >> "$SMODES"  # 0100755
echo "sif /usr/libexec                    mode 16877" >> "$SMODES"  # 0040755
echo "sif /home/root/.ssh                 mode 16832" >> "$SMODES"  # 0040700
echo "sif /home/root/.ssh/authorized_keys mode 33152" >> "$SMODES"  # 0100600
DBG2="$WORK/dbg2.log"
debugfs -w -f "$SMODES" "$WORK/p2" > "$DBG2" 2>&1 || err "debugfs mode-setting failed"
if grep -qiE 'error|failed|denied|not found' "$DBG2"; then
  sed 's/^/    /' "$DBG2" >&2
  err "debugfs reported errors (log above)"
fi

say "post-injection filesystem check (e2fsck -fn) ..."
e2fsck -fn "$WORK/p2" >/dev/null 2>&1 || err "post-injection e2fsck failed — p2 not written back; output image is an unmodified copy"

# sanity: the injected objects must be in place with the right modes and 0:0
check_obj() {
  debugfs -R "stat $1" "$WORK/p2" 2>/dev/null | grep -q "Mode: *$2" || err "injected $1 missing or wrong mode (want $2)"
  debugfs -R "stat $1" "$WORK/p2" 2>/dev/null | grep -q 'User:     0   Group:     0' || err "injected $1 not owned by 0:0"
}
check_obj /usr/local/bin/dropbear         0755
check_obj /usr/local/bin/dropbearkey      0755
check_obj /usr/libexec/sftp-server        0755
check_obj /etc/init.d/sshd                0755
check_obj /home/root/.ssh                 0700
check_obj /home/root/.ssh/authorized_keys 0600
debugfs -R 'stat /etc/rc5.d/S30sshd' "$WORK/p2" 2>/dev/null | grep -q 'Fast link dest: "../init.d/sshd"' \
  || err "boot symlink /etc/rc5.d/S30sshd missing or wrong target"

# binaries in the p2 must match the ones we injected
for B in dropbear dropbearkey sftp-server; do
  case "$B" in
    dropbear)      P=/usr/local/bin/dropbear ;;
    dropbearkey)   P=/usr/local/bin/dropbearkey ;;
    sftp-server)   P=/usr/libexec/sftp-server ;;
  esac
  debugfs -R "dump $P $WORK/check-$B" "$WORK/p2" 2>/dev/null
  [ "$(sha256 "$WORK/check-$B")" = "$(sha256 "$BINDIR/$B")" ] || err "injected $B does not match $BINDIR/$B"
  rm -f "$WORK/check-$B"
done

say "writing p2 back at offset $P2_OFFSET ..."
dd if="$WORK/p2" of="$OUT_REAL" bs=512 seek="$P2_SECTORS" conv=notrunc 2>/dev/null

# --- verification ---------------------------------------------------------------
say "verifying ..."
OUT_SIZE=$(wc -c < "$OUT_REAL" | tr -d '[:space:]')
[ "$OUT_SIZE" = "$IN_SIZE" ] || err "output size changed: $OUT_SIZE != $IN_SIZE"

# every byte before and after p2 must be identical to the input
say "comparing bytes before p2 (0 .. $P2_OFFSET) ..."
cmp -n "$P2_OFFSET" "$IN_REAL" "$OUT_REAL" || err "bytes before p2 differ — output discarded"
say "comparing bytes after p2 ($((P2_OFFSET + P2_SIZE)) .. end) ..."
cmp -i "$((P2_OFFSET + P2_SIZE))" "$IN_REAL" "$OUT_REAL" || err "bytes after p2 differ (game/data partitions!) — output discarded"

# input must be untouched
IN_SIZE2=$(wc -c < "$IN_REAL" | tr -d '[:space:]')
[ "$IN_SIZE2" = "$IN_SIZE" ] || err "input image size changed — this must not happen"

say "done."
say "output image : $OUT_REAL"
say "  sha256     : $(sha256 "$OUT_REAL")"
say ""
say "next steps:"
say "  - flash to a full-size SD card and insert into the machine's mmcblk0 slot"
say "  - WARNING: the image's /data (scores/settings) is factory-empty; flashing over"
say "    a card in use erases it. Back up the live card's data first (see README)."
say "  - on first boot the machine generates Dropbear host keys into /data/dropbear"
say "  - then: ssh -i <private-key> root@<machine-ip>   (key-only, port 22)"
say "  - optional deeper check before flashing:"
say "      ./validate-image.sh -i $OUT_REAL -k $KEYFILE"