#!/bin/sh
# Install the Fillrate bare-metal runner once.
# New tasks do not need a new install. The runner only executes a task that
# Fillrate has signed, plus a one-time ticket for this machine and time window.
#
#   curl -fsSL https://fillrate.net/install/bare-metal.sh | sudo bash -s -- \
#     --machine-id 146922 --public-key 'ssh-ed25519 AAAA...'
#
# View the script before you run it: https://fillrate.net/install/bare-metal.sh
#
# Read-only diagnostics access (skip with --skip-diag-access): the fillrate
# user joins adm and systemd-journal so it can read kernel and system logs
# (no extra sudo). rasdaemon is installed and enabled to record memory and
# machine-check errors, an EDAC driver is loaded when the CPU has one, and
# ipmitool is installed only when an IPMI/BMC device is present. The adm
# group can read the rasdaemon event DB. At most one EDAC driver is tried,
# only when the kernel modalias tables match it to this hardware. apt runs only
# when a simulation shows new installs plus revision-only dependency upgrades
# (same upstream version, only the Debian/Ubuntu revision changes), no
# removals, and nothing NVIDIA, CUDA, Docker, containerd, Vast, kernel, core
# libc, or systemd. Nothing is
# restarted and the machine is not rebooted. The same steps exist as the
# signed task diag_access_setup for hosts installed before this change.
set -eu

ORIGIN=https://fillrate.net
RUNNER_SHA256=0232267ac28b479c51efa326f51e899d6890b6d83b96dae10ff81759b739d03e
JUMP_IP=134.199.240.201
MACHINE_ID=""
PUBLIC_KEY=""
TASK_PUB_FILE=""
TICKET_PUB_FILE=""
RUNNER_FILE=""
DIAG_ACCESS=1

die() {
  printf '%s\n' "bare-metal install: $*" >&2
  exit 1
}

while [ $# -gt 0 ]; do
  case "$1" in
    --machine-id)
      MACHINE_ID=${2:-}
      shift 2
      ;;
    --public-key)
      PUBLIC_KEY=${2:-}
      shift 2
      ;;
    --origin)
      ORIGIN=${2:-}
      shift 2
      ;;
    --task-pub-file)
      TASK_PUB_FILE=${2:-}
      shift 2
      ;;
    --ticket-pub-file)
      TICKET_PUB_FILE=${2:-}
      shift 2
      ;;
    --runner-file)
      RUNNER_FILE=${2:-}
      shift 2
      ;;
    --skip-diag-access)
      DIAG_ACCESS=0
      shift
      ;;
    *)
      die "unknown argument: $1"
      ;;
  esac
done

[ "$(id -u)" -eq 0 ] || die "run as root"
printf '%s\n' "$MACHINE_ID" | grep -E -q '^[A-Za-z0-9._:-]{1,64}$' || die "machine-id is invalid"
case "$PUBLIC_KEY" in
  ssh-*) ;;
  *) die "public-key must be an OpenSSH public key line" ;;
esac

tmp=$(mktemp -d)
trap 'rm -rf "$tmp"' EXIT

if [ -n "$RUNNER_FILE" ]; then
  cp -- "$RUNNER_FILE" "$tmp/fillrate-runner"
else
  curl -fsSL "$ORIGIN/install/fillrate-runner" -o "$tmp/fillrate-runner" || die "could not download fillrate-runner"
fi
printf '%s  %s\n' "$RUNNER_SHA256" "$tmp/fillrate-runner" | sha256sum -c - || die "fillrate-runner checksum did not match"

if [ -n "$TASK_PUB_FILE" ]; then
  task_pub=$(sed -n '1p' "$TASK_PUB_FILE" | tr -d '\r')
else
  task_pub=""
fi
if [ -n "$TICKET_PUB_FILE" ]; then
  ticket_pub=$(sed -n '1p' "$TICKET_PUB_FILE" | tr -d '\r')
else
  ticket_pub=""
fi
if [ -z "$task_pub" ] || [ -z "$ticket_pub" ]; then
  keys=$(curl -fsSL "$ORIGIN/install/signing-keys") || die "could not download signing keys"
  [ -n "$task_pub" ] || task_pub=$(printf '%s\n' "$keys" | sed -n 's/^task //p' | sed -n '1p')
  [ -n "$ticket_pub" ] || ticket_pub=$(printf '%s\n' "$keys" | sed -n 's/^ticket //p' | sed -n '1p')
fi
case "$task_pub" in
  ssh-ed25519\ *) ;;
  *) die "task signing key is missing" ;;
esac
case "$ticket_pub" in
  ssh-ed25519\ *) ;;
  *) die "ticket signing key is missing" ;;
esac

printf '%s\n' "$PUBLIC_KEY" > "$tmp/host.pub"
ssh-keygen -l -f "$tmp/host.pub" >/dev/null 2>&1 || die "public key was rejected by ssh-keygen"
printf '%s\n' "$task_pub" > "$tmp/task.pub"
ssh-keygen -l -f "$tmp/task.pub" >/dev/null 2>&1 || die "task signing key was rejected by ssh-keygen"
printf '%s\n' "$ticket_pub" > "$tmp/ticket.pub"
ssh-keygen -l -f "$tmp/ticket.pub" >/dev/null 2>&1 || die "ticket signing key was rejected by ssh-keygen"

install -d -m 755 -o root -g root /etc/fillrate
printf '%s\n' "$MACHINE_ID" > /etc/fillrate/machine_id
printf '%s\n' "$task_pub" > /etc/fillrate/task-signing.pub
printf '%s\n' "$ticket_pub" > /etc/fillrate/ticket-signing.pub
chmod 644 /etc/fillrate/machine_id /etc/fillrate/task-signing.pub /etc/fillrate/ticket-signing.pub
chown root:root /etc/fillrate/machine_id /etc/fillrate/task-signing.pub /etc/fillrate/ticket-signing.pub

install -m 755 -o root -g root "$tmp/fillrate-runner" /usr/local/sbin/fillrate-runner

sudoers=$tmp/fillrate.sudoers
printf '%s\n' 'fillrate ALL=(root) NOPASSWD: /usr/local/sbin/fillrate-runner' > "$sudoers"
visudo -cf "$sudoers" >/dev/null || die "sudoers file failed visudo"
install -m 440 -o root -g root "$sudoers" /etc/sudoers.d/fillrate

if ! id fillrate >/dev/null 2>&1; then
  useradd --system --create-home --home-dir /var/lib/fillrate --shell /bin/bash --comment "Fillrate bare-metal runner" fillrate
fi
install -d -m 700 -o fillrate -g fillrate /var/lib/fillrate/.ssh
printf 'from="%s" %s\n' "$JUMP_IP" "$PUBLIC_KEY" > /var/lib/fillrate/.ssh/authorized_keys
chown fillrate:fillrate /var/lib/fillrate/.ssh/authorized_keys
chmod 600 /var/lib/fillrate/.ssh/authorized_keys
install -d -m 700 -o root -g root /var/lib/fillrate-runner

# plan_check REQUESTED... < apt-get -s output
# Allows new installs outside the protected set, plus revision-only upgrades
# of dependencies: same epoch and upstream version, only the Debian/Ubuntu
# revision after the last '-' changes, not a requested package, not protected.
# Refuses any removal and anything else. Prints one line per decision.
plan_check() {
  awk -v req=" $* " '
    function prot(p) {
      return p ~ /^(nvidia|libnvidia|cuda|libcuda|docker|containerd|vast|linux-image|linux-modules|linux-headers|libc6|libc-bin|libc-dev-bin|libc-l10n|locales|systemd|libsystemd|udev|libudev)/
    }
    function epoch(v) { return (index(v, ":") ? substr(v, 1, index(v, ":") - 1) : "0") }
    function noep(v) { return (index(v, ":") ? substr(v, index(v, ":") + 1) : v) }
    function upstream(v,   s, i) { s = noep(v); i = match(s, /-[^-]*$/); return (i ? substr(s, 1, i - 1) : "") }
    /^(Remv|Purg) / { print "deny: removes " $2; bad = 1; next }
    /^Inst / {
      p = $2
      if (prot(p)) { print "deny: protected package " p; bad = 1; next }
      if ($3 ~ /^\[/) {
        old = $3; gsub(/[\[\]]/, "", old)
        new = $4; gsub(/[()]/, "", new)
        if (index(req, " " p " ")) { print "deny: requested package " p " is an upgrade"; bad = 1; next }
        uo = upstream(old); un = upstream(new)
        if (uo == "" || un == "" || uo != un || epoch(old) != epoch(new)) {
          print "deny: upgrade " p " " old " -> " new " changes more than the revision"; bad = 1; next
        }
        print "allow-upgrade: " p " " old " -> " new " (revision-only dependency upgrade)"
        next
      }
      print "allow-install: " p
    }
    END { exit bad }
  '
}

EDAC_SYS=/sys/devices/system/edac/mc
SYSFS=/sys
PROC_MODULES=/proc/modules
MODULES_LOAD=/etc/modules-load.d/fillrate-edac.conf
RAS_TMPFILES=/etc/tmpfiles.d/fillrate-rasdaemon.conf
RAS_DROPIN=/etc/systemd/system/rasdaemon.service.d/fillrate-perms.conf
EDAC_MODE=apply
EDAC_UNLOADED=""
RAS_PERMS=skipped
edac_say() { printf '%s\n' "diag access: $*"; }
edac_act() { "$@" >/dev/null 2>&1; }
# ---- EDAC helpers (shared with install/bare-metal.sh) ----
# A driver is a candidate only when the kernel's own alias tables match it to
# this machine: modprobe -R on the CPU modalias and on each PCI modalias,
# keeping *_edac drivers (not *_edac_common). At most one driver is ever
# tried. Dry-run and apply run the same decisions; only edac_act differs.
edac_bound() { ls -d "$EDAC_SYS"/mc[0-9]* >/dev/null 2>&1; }
edac_loaded() { awk '{print $1}' "$PROC_MODULES" 2>/dev/null; }
edac_is_loaded() { edac_loaded | grep -x -q "$1"; }
edac_candidates() {
  {
    cat "$SYSFS/devices/system/cpu/modalias" 2>/dev/null
    cat "$SYSFS"/bus/pci/devices/*/modalias 2>/dev/null
  } | while IFS= read -r a; do
    [ -n "$a" ] && modprobe -R "$a" 2>/dev/null
  done | grep -E '_edac$' | grep -v -E '_edac_common$' | sort -u
}
# Holders of module $1 that are not in the space-separated set $2.
edac_other_holders() {
  for h in "$SYSFS/module/$1/holders"/*; do
    [ -e "$h" ] || continue
    n=$(basename "$h")
    case " $2 " in *" $n "*) ;; *) printf '%s\n' "$n" ;; esac
  done
}
edac_setup() {
  if edac_bound; then
    EDAC_RESULT=already-loaded
    edac_say "EDAC memory controller already registered"
    return 0
  fi
  cands=$(edac_candidates | tr '\n' ' ' | sed 's/ *$//')
  edac_say "EDAC candidates from modalias: ${cands:-none}"
  # Stray cleanup: nothing is bound, so any loaded *_edac driver the hardware
  # does not match (for example igen6_edac on a server CPU) is unused.
  unload=""
  for m in $(edac_loaded | grep -E '_edac$' | grep -v -E '_edac_common$'); do
    case " $cands " in *" $m "*) continue ;; esac
    [ -z "$(edac_other_holders "$m" "")" ] && unload="$unload $m"
  done
  for m in $(edac_loaded | grep -E '_edac_common$'); do
    [ -z "$(edac_other_holders "$m" "$unload")" ] && unload="$unload $m"
  done
  for m in $unload; do
    if edac_act rmmod "$m"; then
      edac_say "unloaded unbound EDAC module $m"
      EDAC_UNLOADED="$EDAC_UNLOADED $m"
      if [ -f "$MODULES_LOAD" ] && grep -x -q "$m" "$MODULES_LOAD"; then edac_act rm -f "$MODULES_LOAD"; fi
    fi
  done
  set -- $cands
  if [ $# -eq 0 ]; then
    EDAC_RESULT=none-available
    edac_say "no EDAC driver matches this machine; memory errors rely on rasdaemon MCE records"
    return 0
  fi
  if [ $# -gt 1 ]; then
    EDAC_RESULT=none-available
    edac_say "more than one EDAC driver matches ($cands); not loading any"
    return 0
  fi
  m=$1
  if edac_is_loaded "$m"; then
    EDAC_RESULT=none-available
    edac_say "$m is loaded but registered no memory controller (firmware-first or EDAC off in BIOS)"
    return 0
  fi
  before=$(edac_loaded | sort)
  if [ "$EDAC_MODE" = dry-run ]; then
    edac_say "would run: modprobe $m (once; unloaded again if no memory controller registers)"
    EDAC_RESULT=would-try:$m
    return 0
  fi
  edac_act modprobe "$m"
  if edac_bound; then
    EDAC_RESULT=loaded:$m
    edac_act sh -c "printf '%s\n' '# Added by Fillrate diag access setup' '$m' > '$MODULES_LOAD' && chmod 644 '$MODULES_LOAD'"
    edac_say "persisted $m in $MODULES_LOAD"
    return 0
  fi
  # Nothing bound: remove whatever this attempt loaded, newest first.
  bf=$(mktemp) || bf=/dev/null
  printf '%s\n' "$before" > "$bf"
  new=$(edac_loaded | sort | comm -13 "$bf" -)
  rm -f "$bf"
  # Two passes so a dependency is removed after the module that held it.
  for pass in 1 2; do
    for n in $new; do
      edac_is_loaded "$n" || continue
      [ -z "$(edac_other_holders "$n" "")" ] && edac_act rmmod "$n" && EDAC_UNLOADED="$EDAC_UNLOADED $n"
    done
  done
  EDAC_RESULT=none-available
  edac_say "$m did not register a memory controller; unloaded what this attempt loaded"
  return 0
}

# ---- rasdaemon DB read access for the adm group ----
# Directory 0750 root:adm and the event DB 0640 root:adm, applied now and by
# systemd-tmpfiles at boot and after each rasdaemon start (drop-in), so a
# recreated DB gets the same mode. No service is restarted.
RAS_TMPFILES_BODY='d /var/lib/rasdaemon 0750 root adm -
z /var/lib/rasdaemon/ras-mc_event.db 0640 root adm -
z /var/lib/rasdaemon/ras-mc_event.db-journal 0640 root adm -'
RAS_DROPIN_BODY='[Service]
ExecStartPost=-/bin/sh -c "i=0; while [ $i -lt 15 ] && [ ! -e /var/lib/rasdaemon/ras-mc_event.db ]; do i=$((i+1)); sleep 1; done; /bin/systemd-tmpfiles --create /etc/tmpfiles.d/fillrate-rasdaemon.conf"'
ras_perms_setup() {
  changed=0
  if [ "$(cat "$RAS_TMPFILES" 2>/dev/null)" != "$RAS_TMPFILES_BODY" ]; then
    edac_act sh -c "mkdir -p '$(dirname "$RAS_TMPFILES")' && cat > '$RAS_TMPFILES'" << TF || return 1
$RAS_TMPFILES_BODY
TF
    changed=1
  fi
  if [ "$(cat "$RAS_DROPIN" 2>/dev/null)" != "$RAS_DROPIN_BODY" ]; then
    edac_act sh -c "mkdir -p '$(dirname "$RAS_DROPIN")' && cat > '$RAS_DROPIN'" << DI || return 1
$RAS_DROPIN_BODY
DI
    edac_act systemctl daemon-reload
    changed=1
  fi
  edac_act systemd-tmpfiles --create "$RAS_TMPFILES" || return 1
  [ $changed = 1 ] && RAS_PERMS=set || RAS_PERMS=already-set
  return 0
}

diag_access_setup() {
  export DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l NEEDRESTART_SUSPEND=1
  for g in adm systemd-journal; do
    if getent group "$g" >/dev/null 2>&1; then
      usermod -aG "$g" fillrate && printf '%s\n' "diag access: fillrate added to $g"
    fi
  done
  command -v apt-get >/dev/null 2>&1 || { printf '%s\n' "diag access: no apt-get, skipped packages"; return 0; }
  pkgs=""
  dpkg-query -W -f='${Status}' rasdaemon 2>/dev/null | grep -q 'install ok installed' || pkgs="rasdaemon"
  if [ -e /dev/ipmi0 ] || [ -e /dev/ipmi/0 ] || [ -e /dev/ipmidev/0 ] || \
     { command -v dmidecode >/dev/null 2>&1 && dmidecode -t 38 2>/dev/null | grep -q 'IPMI Device Information'; }; then
    dpkg-query -W -f='${Status}' ipmitool 2>/dev/null | grep -q 'install ok installed' || pkgs="$pkgs ipmitool"
  fi
  if [ -n "$pkgs" ]; then
    apt-cache policy $pkgs 2>/dev/null | grep -q 'Candidate: [0-9]' || apt-get update -q >/dev/null 2>&1 || true
    plan=$(apt-get -s install --no-install-recommends $pkgs 2>&1) || { printf '%s\n' "diag access: apt simulation failed, skipped packages"; return 0; }
    if ! verdict=$(printf '%s\n' "$plan" | plan_check $pkgs); then
      printf '%s\n' "$verdict" | sed -n 's/^deny: /diag access: /p'
      printf '%s\n' "diag access: apt plan has a removal, a protected package, or a non-revision upgrade; skipped $pkgs"
    else
      printf '%s\n' "$verdict" | sed -n 's/^allow-upgrade: /diag access: allowed upgrade /p'
      apt-get install -y -q --no-install-recommends -o Dpkg::Options::=--force-confold $pkgs >/dev/null 2>&1 \
        && printf '%s\n' "diag access: installed $pkgs" \
        || printf '%s\n' "diag access: could not install $pkgs"
    fi
  fi
  if dpkg-query -W -f='${Status}' rasdaemon 2>/dev/null | grep -q 'install ok installed'; then
    systemctl enable --now rasdaemon.service >/dev/null 2>&1 || printf '%s\n' "diag access: could not enable rasdaemon"
  fi
  if dpkg-query -W -f='${Status}' rasdaemon 2>/dev/null | grep -q 'install ok installed'; then
    ras_perms_setup && printf '%s\n' "diag access: rasdaemon DB readable by adm ($RAS_PERMS)" || printf '%s\n' "diag access: could not set rasdaemon DB access"
  fi
  edac_setup
  printf '%s\n' "diag access: EDAC $EDAC_RESULT"
}
if [ "$DIAG_ACCESS" = 1 ]; then
  diag_access_setup || printf '%s\n' "diag access: setup did not finish; the runner install is unaffected"
fi

printf '%s\n' "Installed $(/usr/local/sbin/fillrate-runner version)"
printf '%s\n' "SSH user is fillrate. Access is limited to $JUMP_IP."
printf '%s\n' "Remove it with: sudo fillrate-runner uninstall"