#!/usr/bin/bash
# install-headless-guard — wrap every Playwright browser binary so a headless run
# cannot start on the workstation.
#
# WHY A BINARY WRAPPER: Playwright launches the browser by ABSOLUTE PATH out of
# ~/.cache/ms-playwright. No PATH shim and no PreToolUse Bash hook can see that —
# a driver started by an already-running process, an MCP server, or a test runner
# bypasses both. The executable itself is the only place a guarantee exists.
#
# Browsers belong on a buildbox, so the guard blocks only where they do not belong.
# Host identity comes from the buildbox registry: a host the registry lists passes
# through, everything else blocks — including an unreadable or malformed registry.
# E2E_REMOTE_OK=1 stays the workstation's escape hatch for e2e-remote's own payload.
# The resolved host's machine_id is baked into a pass-through wrapper and re-checked
# against /etc/machine-id at launch, so a wrapper written elsewhere cannot unguard
# this machine.
#
# Idempotent: the guard body is the artifact. A wrapper whose body differs from the
# one this script would write now is rewritten; an identical body is left untouched,
# so re-wrapping does not touch the cache mtime that headless-guard.path watches.
set -uo pipefail

PW_DIR="${PLAYWRIGHT_BROWSERS_PATH:-$HOME/.cache/ms-playwright}"
MARKER_FAMILY="e2e-headless-guard"
MARKER="$MARKER_FAMILY v2"
REGISTRY_CLI="$(dirname -- "$(readlink -f -- "$0")")/../lib/buildbox-registry.mjs"

usage() {
  cat >&2 <<'USAGE'
usage: install-headless-guard [--uninstall] [--status] [--settle]
                             [--install-units] [--dir <ms-playwright dir>]

Wraps headless-capable Playwright browsers (chromium headless shell, firefox,
webkit) so they refuse to run on a host the buildbox registry does not list.
Headed chromium-*/chrome-linux64/chrome is deliberately NOT wrapped: a visible
browser is interactive debugging, not the automated fleet that saturates the
machine.

  --uninstall      restore every wrapped binary from its .real sibling
  --status         report wrapped/stale/unwrapped counts and the resolved host
                   identity, exit 1 if any wrapper is missing or out of date
  --settle         wrap repeatedly until the browser cache stops changing
                   (a download creates its version dir before extracting the
                   binary into it, so one pass can arrive too early)
  --install-units  link + enable the systemd --user path/timer that run --settle
USAGE
  exit "${1:-2}"
}

RUN_MODE=install
while [ $# -gt 0 ]; do
  case "$1" in
    --uninstall) RUN_MODE=uninstall; shift ;;
    --status)    RUN_MODE=status; shift ;;
    --settle)    RUN_MODE=settle; shift ;;
    --install-units) RUN_MODE=units; shift ;;
    --dir)       PW_DIR="${2-}"; shift 2 ;;
    -h|--help)   usage 0 ;;
    *)           echo "install-headless-guard: unknown arg: $1" >&2; usage ;;
  esac
done

if [ "$RUN_MODE" = units ]; then
  src="$HOME/.claude/systemd/user"
  dst="$HOME/.config/systemd/user"
  mkdir -p "$dst" || exit 1
  for u in headless-guard.service headless-guard.path headless-guard.timer; do
    [ -f "$src/$u" ] || { echo "install-headless-guard: FATAL missing unit $src/$u" >&2; exit 1; }
    ln -sfn "$src/$u" "$dst/$u" || exit 1
  done
  systemctl --user daemon-reload || exit 1
  systemctl --user enable --now headless-guard.path headless-guard.timer >/dev/null || exit 1
  echo "install-headless-guard: units linked from $src and enabled"
  exit 0
fi

# Identity, resolved once per run. SELF_MACHINE_ID non-empty is the only thing that
# makes a wrapper pass through, so every failure to resolve leaves the guard blocking.
SELF_MACHINE_ID=""
IDENTITY="unresolved"
resolve_identity() {
  local out rc node_bin name machine_id via
  node_bin="$(command -v node 2>/dev/null)" || node_bin=""
  if [ -z "$node_bin" ]; then IDENTITY="unresolved (no node)"; return; fi
  if [ ! -f "$REGISTRY_CLI" ]; then IDENTITY="unresolved (no registry reader at $REGISTRY_CLI)"; return; fi
  out="$("$node_bin" "$REGISTRY_CLI" self 2>/dev/null)"; rc=$?
  case "$rc" in
    0)
      read -r name machine_id via <<<"$out"
      if [ -z "${machine_id:-}" ]; then IDENTITY="unresolved (registry reader returned no machine_id)"; return; fi
      SELF_MACHINE_ID="$machine_id"
      IDENTITY="buildbox $name (via $via)" ;;
    4) IDENTITY="workstation" ;;
    *) IDENTITY="unresolved (registry unreadable)" ;;
  esac
}
resolve_identity
if [ -n "$SELF_MACHINE_ID" ]; then GUARD_MODE="pass-through"; else GUARD_MODE="blocking"; fi

[ -d "$PW_DIR" ] || { echo "install-headless-guard: no browser dir at $PW_DIR — nothing to do (identity=$IDENTITY)"; exit 0; }

if [ "$RUN_MODE" = settle ]; then
  prev=""; stable=0
  for _ in $(seq 1 120); do
    "$0" --dir "$PW_DIR" || exit 1
    cur="$(find "$PW_DIR" -maxdepth 4 -printf '%T@ %s %p\n' 2>/dev/null | sha256sum)"
    if [ "$cur" = "$prev" ]; then stable=$((stable + 1)); else stable=0; fi
    if [ "$stable" -ge 2 ]; then
      echo "install-headless-guard: cache stable, settled"
      exit 0
    fi
    prev="$cur"
    sleep 15
  done
  echo "install-headless-guard: cache still changing after 30min; everything present is wrapped" >&2
  exit 0
fi

# Headless-capable engines only. Headed chrome is left alone on purpose (see usage).
targets() {
  find "$PW_DIR" -maxdepth 4 -type f \
    \( -name 'chrome-headless-shell' -o -path '*/firefox-*/firefox/firefox' -o -name 'pw_run.sh' \) \
    2>/dev/null | sort
}

guard_body() {
  local real=$1 name=$2 passthrough=""
  if [ -n "$SELF_MACHINE_ID" ]; then
    passthrough="if [ \"\$(cat /etc/machine-id 2>/dev/null)\" = \"$SELF_MACHINE_ID\" ]; then
  exec \"$real\" \"\$@\"
fi"
  fi
  cat <<GUARD
#!/usr/bin/bash
# $MARKER — DO NOT EDIT. Reinstall with ~/.claude/bin/install-headless-guard
if [ "\${E2E_REMOTE_OK:-}" = "1" ]; then
  exec "$real" "\$@"
fi
$passthrough
cat >&2 <<'MSG'
BLOCKED: $name may not run on this machine.

A headless browser and its dev server are one coupled pair and belong on
a buildbox — run on the workstation they pin ~250% CPU to the laptop. Launch the
whole pair with the wrapper, which starts BOTH halves remotely and rsyncs
artifacts back:

  ~/.claude/bin/e2e-remote --server "<dev-server> --host 127.0.0.1 --port <P>" \\
    --wait-port <P> --env PORT=<P> --mkdir <outdir> -- <browser-client>

This host is not a buildbox: no host in ~/.claude/buildbox-hosts.json matches its
identity. If it IS one, fix its registry entry and rerun install-headless-guard
instead of exporting the override.

Deliberate local run (boxes down, headed debugging)? Export E2E_REMOTE_OK=1.
If this fired right after a browser version bump, rerun install-headless-guard.
MSG
exit 97
GUARD
}

write_guard() {
  local dest=$1 real=$2 name=$3
  guard_body "$real" "$name" >"$dest" || return 1
  chmod 755 "$dest"
}

# Bounded to the header the guard writes: an unwrapped target is a stripped ELF, and
# scanning it whole would both cost a full read per pass and let payload bytes pose as
# the marker. Process substitution, not a pipe, so the status is grep's alone.
is_guard() { grep -qF "$MARKER_FAMILY" <(head -c 200 -- "$1" 2>/dev/null); }
guard_is_current() { [ "$(cat -- "$1")" = "$(guard_body "$2" "$(basename -- "$1")")" ]; }

# Rewrites only a body that differs, and only through a same-directory rename: a
# no-op pass must not touch the cache that headless-guard.path and --settle watch.
# 0 = body replaced, 1 = already current, 2 = failed.
converge_guard() {
  local bin=$1 real=$2 tmp="$1.guard.$$"
  guard_is_current "$bin" "$real" && return 1
  write_guard "$tmp" "$real" "$(basename -- "$bin")" || { rm -f -- "$tmp"; return 2; }
  mv -f -- "$tmp" "$bin" || { rm -f -- "$tmp"; return 2; }
  return 0
}

wrapped=0; refreshed=0; current=0; restored=0; unwrapped=0; stale=0
while IFS= read -r bin; do
  [ -n "$bin" ] || continue
  real="$bin.real"
  case "$RUN_MODE" in
    install)
      if [ -e "$real" ]; then
        is_guard "$bin" || {
          echo "install-headless-guard: FATAL $real exists but $bin is not a guard" >&2; exit 1; }
        converge_guard "$bin" "$real"
        case $? in
          0) refreshed=$((refreshed + 1)) ;;
          1) current=$((current + 1)) ;;
          *) echo "install-headless-guard: FATAL cannot rewrite guard at $bin" >&2; exit 1 ;;
        esac
        continue
      fi
      # Moving a guard onto .real would destroy the browser it is standing in for.
      is_guard "$bin" && {
        echo "install-headless-guard: FATAL $bin is a guard with no .real sibling" >&2; exit 1; }
      mv -- "$bin" "$real" || { echo "install-headless-guard: FATAL cannot move $bin" >&2; exit 1; }
      if ! write_guard "$bin" "$real" "$(basename "$bin")"; then
        mv -- "$real" "$bin"                      # restore rather than leave it unlaunchable
        echo "install-headless-guard: FATAL cannot write guard at $bin" >&2
        exit 1
      fi
      wrapped=$((wrapped + 1)) ;;
    uninstall)
      [ -e "$real" ] || continue
      mv -f -- "$real" "$bin" && restored=$((restored + 1)) ;;
    status)
      if [ -e "$real" ] && is_guard "$bin"; then
        if guard_is_current "$bin" "$real"; then
          wrapped=$((wrapped + 1))
        else
          stale=$((stale + 1)); echo "STALE $bin"
        fi
      else
        unwrapped=$((unwrapped + 1)); echo "UNWRAPPED $bin"
      fi ;;
  esac
done <<EOF
$(targets)
EOF

case "$RUN_MODE" in
  install)   echo "install-headless-guard: wrapped=$wrapped refreshed=$refreshed current=$current mode=$GUARD_MODE identity=$IDENTITY ($PW_DIR)" ;;
  uninstall) echo "install-headless-guard: restored=$restored ($PW_DIR)" ;;
  status)
    echo "install-headless-guard: wrapped=$wrapped stale=$stale unwrapped=$unwrapped mode=$GUARD_MODE identity=$IDENTITY ($PW_DIR)"
    [ "$unwrapped" -eq 0 ] && [ "$stale" -eq 0 ] || exit 1 ;;
esac
