#!/usr/bin/env bash
# Runs AS ROOT on a build box. Holds a boot-affecting change hostage: last-known-good
# is restored unless a session arrives under a boot id different from the one that was
# current when the change was applied. A change that is fine live and fatal at the next
# cold boot is the failure this exists to survive, so the timer is boot-persistent and
# the cancel predicate is a changed boot id — never a fresh connection, which the
# bricking change on debian2 would also have satisfied.
set -euo pipefail

STATE_DIR="${BUILDBOX_DEADMAN_STATE_DIR:-/var/lib/buildbox-deadman}"
STATE="$STATE_DIR/state.json"
UNIT_DIR="${BUILDBOX_DEADMAN_UNIT_DIR:-/etc/systemd/system}"
SELF="${BUILDBOX_DEADMAN_SELF:-/usr/local/sbin/buildbox-deadman}"
# Kept empty on a box.  The prefix lets the branch tests exercise snapshot and restore
# decisions against a disposable filesystem, rather than a machine's /etc.
ROOT="${BUILDBOX_DEADMAN_ROOT:-}"
GUARDED=(etc/ssh etc/systemd etc/sysctl.d etc/fstab)

boot_id() {
  if [ -n "${BUILDBOX_DEADMAN_BOOT_ID_FILE:-}" ]; then tr -d '\n' <"$BUILDBOX_DEADMAN_BOOT_ID_FILE"
  else tr -d '\n' </proc/sys/kernel/random/boot_id; fi
}
now() { printf '%s\n' "${BUILDBOX_DEADMAN_NOW:-$(date +%s)}"; }
die() { printf 'buildbox-deadman: %s\n' "$*" >&2; exit 1; }

# Test-only override keeps the branch test hermetic on an unprivileged workstation;
# production always derives this from the kernel identity.
EFFECTIVE_UID="${BUILDBOX_DEADMAN_EFFECTIVE_UID:-$(id -u)}"
[ "$EFFECTIVE_UID" = 0 ] || die "must run as root"

state_get() { python3 -c 'import json,sys; print(json.load(open(sys.argv[1])).get(sys.argv[2],""))' "$STATE" "$1"; }

install_units() {
  cat >"$UNIT_DIR/buildbox-deadman.service" <<'UNIT'
[Unit]
Description=Restore last-known-good host config unless a post-reboot session cancelled it
DefaultDependencies=no
After=local-fs.target
[Service]
Type=oneshot
ExecStart=/usr/local/sbin/buildbox-deadman check
UNIT
  cat >"$UNIT_DIR/buildbox-deadman.timer" <<'UNIT'
[Unit]
Description=Deadman check for a boot-affecting host-config change
[Timer]
OnBootSec=30s
OnUnitActiveSec=30s
AccuracySec=1s
[Install]
WantedBy=timers.target
UNIT
  systemctl daemon-reload
  systemctl enable --now buildbox-deadman.timer >/dev/null
}

snapshot() {
  local out="$1" paths=()
  for p in "${GUARDED[@]}"; do [ -e "$ROOT/$p" ] && paths+=("$p"); done
  tar -C "${ROOT:-/}" --numeric-owner --acls --xattrs -cf "$out" "${paths[@]}"
}

# A door this host serves right now, checked locally. The port is read from the running
# daemon rather than named here, so the predicate holds on a box whose door moved. Not a
# tailnet probe: a tailscaled hiccup is not a reason to revert a host config, and
# tailscaled's own port 22 listener is not sshd and does not count.
doors_ok() {
  if [ -n "${BUILDBOX_DEADMAN_DOORS_OK:-}" ]; then
    [ "$BUILDBOX_DEADMAN_DOORS_OK" = yes ]
    return
  fi
  systemctl is-active --quiet ssh && ss -H -lntp 2>/dev/null | grep -q '"sshd"'
}

# A snapshot is only worth restoring if the state it captured was itself serving a door.
# The verdict is taken at arm time, when that is observable, and carried in the state
# file — asking the question at restore time would be asking it of a host that is by
# then already broken.
snapshot_verdict() { doors_ok && echo healthy || echo unhealthy; }

# /etc/systemd is guarded, so a revert to a pre-harden snapshot takes the self-healer's
# units with it and nothing converges the door afterward. The script itself lives outside
# the guarded paths and survives, so it is run to put the door back and to say whether
# the box is actually reachable after the revert.
reconverge_door() {
  local access="${BUILDBOX_SSHD_ACCESS:-/usr/local/sbin/buildbox-sshd-access}"
  [ -x "$access" ] && "$access" >&2 || true
  if doors_ok; then
    logger -t buildbox-deadman "door is served after the revert"
    return 0
  fi
  logger -t buildbox-deadman "NO DOOR after the revert — this host needs out-of-band rescue"
  return 1
}

# A truncated or empty archive would extract nothing, and the swap below treats an absent
# path as "delete it" — that turns a corrupt snapshot into a host with no /etc/ssh at all.
snapshot_readable() {
  local snap="$1" listed
  listed="$(tar -tf "$snap" 2>/dev/null)" || return 1
  local p
  for p in "${GUARDED[@]}"; do
    printf '%s\n' "$listed" | grep -qx -e "$p" -e "$p/" && return 0
  done
  return 1
}

restore() {
  local snap="$1" stage="$STATE_DIR/restore.$$"
  rm -rf "$stage"; mkdir -p "$stage"
  tar --numeric-owner --acls --xattrs -xf "$snap" -C "$stage"
  # swap by rename, never edit in place: a crash mid-restore must not leave a host
  # with half an /etc/ssh. A path absent from the snapshot is removed, because the
  # change class includes adding a file that breaks boot.
  for p in "${GUARDED[@]}"; do
    local new="$stage/$p" cur="$ROOT/$p" old="$ROOT/${p}.deadman-old.$$"
    [ -e "$new" ] || { rm -rf "$cur"; continue; }
    if [ -e "$cur" ]; then mv "$cur" "$old"; fi
    mv "$new" "$cur"
    rm -rf "$old"
  done
  rm -rf "$stage"
  systemctl daemon-reload || true
  systemctl restart ssh || systemctl restart sshd || true
}

cmd_arm() {
  local id="" deadline_sec=900
  while [ $# -gt 0 ]; do
    case "$1" in
      --id) id="${2-}"; shift 2 ;;
      --deadline-sec) deadline_sec="${2-}"; shift 2 ;;
      *) die "arm: unknown arg $1" ;;
    esac
  done
  [[ "$id" =~ ^[A-Za-z0-9._-]{1,128}$ ]] || die "arm: --id must match ^[A-Za-z0-9._-]{1,128}$"
  [[ "$deadline_sec" =~ ^[0-9]+$ ]] && [ "$deadline_sec" -ge 60 ] || die "arm: --deadline-sec must be >= 60"
  [ ! -f "$STATE" ] || die "already armed for $(state_get id) — disarm or let it expire before applying another change"

  mkdir -p "$STATE_DIR"; chmod 0700 "$STATE_DIR"
  local snap="$STATE_DIR/lkg-$id.tar"
  snapshot "$snap"
  [ "$(readlink -f "$0")" = "$SELF" ] || install -m 0755 "$0" "$SELF"
  python3 - "$STATE" "$id" "$(boot_id)" "$(( $(now) + deadline_sec ))" "$snap" "$(snapshot_verdict)" <<'PY'
import json,sys
p,i,b,d,s,v = sys.argv[1:7]
json.dump({"id":i,"armedBootId":b,"deadline":int(d),"snapshot":s,"snapshotHealth":v}, open(p,"w"))
PY
  chmod 0600 "$STATE"
  install_units
  echo "DEADMAN armed id=$id bootId=$(boot_id) deadline=+${deadline_sec}s snapshot=$snap health=$(state_get snapshotHealth)"
}

cmd_check() {
  [ -f "$STATE" ] || { systemctl disable --now buildbox-deadman.timer >/dev/null 2>&1 || true; exit 0; }
  local deadline snap id health
  deadline="$(state_get deadline)"; snap="$(state_get snapshot)"; id="$(state_get id)"
  health="$(state_get snapshotHealth)"
  [ "$(now)" -ge "$deadline" ] || exit 0

  # The change is proven the moment this host is up under a different boot id and still
  # serving a door — that is the whole question the deadline was waiting on. Reverting a
  # host that already answered it undoes a repair that worked.
  if [ "$(boot_id)" != "$(state_get armedBootId)" ] && doors_ok; then
    logger -t buildbox-deadman "deadline passed for $id but the host rebooted and still serves its door — cancelling"
    rm -f "$snap" "$STATE"
    systemctl disable --now buildbox-deadman.timer >/dev/null 2>&1 || true
    echo "DEADMAN cancelled id=$id reason=survived-reboot"
    exit 0
  fi

  # A snapshot taken while the host was already broken restores the breakage faithfully.
  # Putting the door back is strictly better than reinstating a state known not to serve one.
  if [ "$health" != healthy ]; then
    logger -t buildbox-deadman "deadline passed for $id but its snapshot was taken from an unhealthy host ($health) — not restoring it"
    mv "$STATE" "$STATE_DIR/refused-$id.json"
    systemctl disable --now buildbox-deadman.timer >/dev/null 2>&1 || true
    reconverge_door || true
    echo "DEADMAN refused id=$id reason=unhealthy-snapshot"
    exit 0
  fi

  if ! snapshot_readable "$snap"; then
    logger -t buildbox-deadman "deadline passed for $id but its snapshot does not list a guarded path — not restoring it"
    mv "$STATE" "$STATE_DIR/refused-$id.json"
    systemctl disable --now buildbox-deadman.timer >/dev/null 2>&1 || true
    reconverge_door || true
    echo "DEADMAN refused id=$id reason=unreadable-snapshot"
    exit 0
  fi

  logger -t buildbox-deadman "deadline passed for $id with no post-reboot cancel — restoring last-known-good"
  restore "$snap"
  mv "$STATE" "$STATE_DIR/reverted-$id.json"
  systemctl disable --now buildbox-deadman.timer >/dev/null 2>&1 || true
  reconverge_door || true
  logger -t buildbox-deadman "restored last-known-good for $id"
  echo "DEADMAN reverted id=$id"
}

cmd_disarm() {
  local id="${1-}"
  [ -f "$STATE" ] || die "not armed"
  [ -z "$id" ] || [ "$id" = "$(state_get id)" ] || die "armed for $(state_get id), not $id"
  local armed current
  armed="$(state_get armedBootId)"; current="$(boot_id)"
  [ "$armed" != "$current" ] || die "refused: boot id is still $current — this host has not rebooted since the change was applied, so nothing has proven it survives a boot"
  rm -f "$STATE_DIR/lkg-$(state_get id).tar" "$STATE"
  systemctl disable --now buildbox-deadman.timer >/dev/null 2>&1 || true
  echo "DEADMAN disarmed bootIdBefore=$armed bootIdAfter=$current"
}

cmd_status() {
  if [ -f "$STATE" ]; then
    echo "DEADMAN armed id=$(state_get id) armedBootId=$(state_get armedBootId) currentBootId=$(boot_id) secondsLeft=$(( $(state_get deadline) - $(now) )) snapshotHealth=$(state_get snapshotHealth)"
  else
    echo "DEADMAN idle currentBootId=$(boot_id)"
  fi
}

case "${1:-}" in
  arm) shift; cmd_arm "$@" ;;
  check) cmd_check ;;
  disarm) shift; cmd_disarm "$@" ;;
  status) cmd_status ;;
  *) echo "usage: buildbox-deadman arm --id <id> [--deadline-sec N] | check | disarm [id] | status" >&2; exit 2 ;;
esac
