#!/bin/sh
# fillrate-task: diag_tools_extended
# Heavier diagnostics that may need a reboot. Run only in an empty, host-
# approved maintenance window.
#   1. install smartmontools, nvme-cli, lm-sensors (new installs only)
#   2. install mcelog only on Intel CPUs and only when apt has a candidate
#      (rasdaemon from diag_access_setup is the main MCE recorder)
#   3. kdump: install kdump-tools, set USE_KDUMP=1, reserve crashkernel=
#      memory, update-grub, enable kdump-tools, then ONE reboot so the
#      reservation takes effect
#   4. after reboot a one-shot unit checks kdump is armed
#      (kdump-config status: "ready to kdump") and writes the report
#
# Safety:
# - Refuses when ANY container is running (rentals must be gone), when
#   apt/dpkg is busy, or when the ticket expiry has passed.
# - apt runs only when apt-get -s passes plan_check (same rule as
#   diag_access_setup): new installs plus logged revision-only dependency
#   upgrades; no removals; nothing protected. Otherwise nothing changes.
# - needrestart is list-only. No docker, vastai, or driver restarts.
# - The reboot happens only with ticket mode=apply, and only when the
#   crashkernel reservation is not already active. If kdump is already
#   armed, there is no reboot.
# - mode=dry-run runs every check and the apt simulation, prints the plan
#   (including whether a reboot would happen), and changes nothing.
# - Idempotent. Every step is logged to stdout and to a mode 644 log under
#   /var/lib/fillrate-runner/diag-tools-extended. No secrets are logged.
# FILLRATE_* path overrides are for tests; a real run uses the defaults.

set -u

STATE=${FILLRATE_DIAGX_STATE:-/var/lib/fillrate-runner/diag-tools-extended}
UNIT=${FILLRATE_DIAGX_UNIT:-/etc/systemd/system/fillrate-diag-tools-verify.service}
KDUMP_DEFAULT=${FILLRATE_KDUMP_DEFAULT:-/etc/default/kdump-tools}
GRUB_CFG=${FILLRATE_GRUB_KDUMP:-/etc/default/grub.d/kdump-tools.cfg}
CMDLINE=${FILLRATE_CMDLINE:-/proc/cmdline}
CPUINFO=${FILLRATE_CPUINFO:-/proc/cpuinfo}
CRASHKERNEL=${FILLRATE_CRASHKERNEL:-crashkernel=0M-32G:512M,32G-:1G}
EXPIRES=${FILLRATE_EXPIRES:-0}
MODE=${FILLRATE_MODE:-}
PHASE=${1:-run}

export DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l NEEDRESTART_SUSPEND=1 APT_LISTCHANGES_FRONTEND=none

umask 022
mkdir -p "$STATE" 2>/dev/null || true
LOG="$STATE/run-$(date -u +%Y%m%dT%H%M%SZ 2>/dev/null || echo now)-$PHASE.log"
: > "$LOG" 2>/dev/null || LOG=/dev/null
chmod 644 "$LOG" 2>/dev/null || true

log() {
  line="$(date -u +%H:%M:%SZ 2>/dev/null) [${MODE:-verify}] $*"
  printf '%s\n' "$line"
  printf '%s\n' "$line" >> "$LOG" 2>/dev/null || true
}
refuse() { log "refusing: $1"; log "summary result=refused"; exit 1; }

kdump_armed() {
  grep -q 'crashkernel=' "$CMDLINE" 2>/dev/null || return 1
  kdump-config status 2>/dev/null | grep -q -i 'ready to kdump'
}

# ---- post-reboot verification (run by the one-shot unit) ----
if [ "$PHASE" = verify ]; then
  log "post-reboot verify"
  if grep -q 'crashkernel=' "$CMDLINE" 2>/dev/null; then log "cmdline has crashkernel"; else log "cmdline missing crashkernel"; fi
  st=$(kdump-config status 2>&1); log "kdump-config status: $st"
  if kdump_armed; then result=armed; else result=not-armed; fi
  for t in smartctl nvme sensors; do command -v "$t" >/dev/null 2>&1 && log "$t present" || log "$t missing"; done
  printf '%s\n' "$result" > "$STATE/kdump-status"
  chmod 644 "$STATE/kdump-status" 2>/dev/null || true
  systemctl disable "$(basename "$UNIT")" >/dev/null 2>&1 || true
  rm -f "$UNIT"
  log "summary result=verified kdump=$result"
  [ "$result" = armed ]
  exit $?
fi

case "$MODE" in
  apply|dry-run) ;;
  *) refuse "mode must be apply or dry-run (the reboot needs an explicit apply)" ;;
esac

# 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 }
  '
}
safe_plan() {
  plan=$(apt-get -s -o Debug::NoLocking=1 install --no-install-recommends "$@" 2>&1) || { log "apt simulation failed"; return 1; }
  verdict=$(printf '%s\n' "$plan" | plan_check "$@")
  vcode=$?
  printf '%s\n' "$verdict" | sed '/^$/d; s/^/    plan: /' | tee -a "$LOG" 2>/dev/null
  return $vcode
}
pkg_installed() { dpkg-query -W -f='${Status}' "$1" 2>/dev/null | grep -q 'install ok installed'; }
act() {
  if [ "$MODE" = dry-run ]; then log "would run: $*"; return 0; fi
  log "run: $*"; "$@" >> "$LOG" 2>&1; c=$?; log "exit=$c: $1"; return $c
}

log "diag_tools_extended start machine=${FILLRATE_MACHINE_ID:-unknown} request=${FILLRATE_BM_REQUEST_ID:-unknown}"
[ "$(id -u)" -eq 0 ] || refuse "not root"
[ "$EXPIRES" = 0 ] || [ "$(date +%s)" -le "$EXPIRES" ] || refuse "ticket expiry has passed"
command -v apt-get >/dev/null 2>&1 || refuse "apt-get not found"
if ! running=$(docker ps -q 2>/dev/null); then refuse "docker ps failed; cannot prove there are no rentals"; fi
[ -z "$running" ] || refuse "containers are running ($(printf '%s\n' "$running" | wc -l | tr -d ' ')); this task needs an empty maintenance window"
if pgrep -x unattended-upgr >/dev/null 2>&1 || pgrep -x apt-get >/dev/null 2>&1; then refuse "apt or unattended-upgrades is running"; fi
if command -v fuser >/dev/null 2>&1 && fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock >/dev/null 2>&1; then refuse "dpkg lock is held"; fi
log "precheck ok: no running containers"

want=""
for p in smartmontools nvme-cli lm-sensors kdump-tools; do pkg_installed "$p" || want="$want $p"; done
vendor=$(awk -F': ' '/^vendor_id/{print $2; exit}' "$CPUINFO" 2>/dev/null)
if [ "$vendor" = GenuineIntel ] && ! pkg_installed mcelog && apt-cache policy mcelog 2>/dev/null | grep -q 'Candidate: [0-9]'; then
  want="$want mcelog"
else
  log "mcelog skipped (vendor=${vendor:-unknown}, or no candidate, or installed)"
fi
want=$(printf '%s' "$want" | sed 's/^ *//')
if [ -n "$want" ]; then
  apt-cache policy $want 2>/dev/null | grep -q 'Candidate: [0-9]' || act apt-get update -q || refuse "apt-get update failed"
  log "simulating: apt-get install --no-install-recommends $want"
  safe_plan $want || refuse "apt plan has a removal, a protected package, or a non-revision upgrade; nothing changed"
  act apt-get install -y -q --no-install-recommends -o Dpkg::Options::=--force-confold $want || refuse "apt install failed"
else
  log "packages already installed"
fi

if kdump_armed; then
  log "kdump already armed; no reboot needed"
  log "summary result=ok reboot=no kdump=armed"
  exit 0
fi

# kdump config
if [ -f "$KDUMP_DEFAULT" ] && grep -q '^USE_KDUMP=1' "$KDUMP_DEFAULT"; then
  log "USE_KDUMP=1 already set"
else
  act sh -c "touch '$KDUMP_DEFAULT' && sed -i '/^USE_KDUMP=/d' '$KDUMP_DEFAULT' && echo USE_KDUMP=1 >> '$KDUMP_DEFAULT'"
fi
if grep -q 'crashkernel=' "$GRUB_CFG" 2>/dev/null; then
  log "crashkernel already in $GRUB_CFG"
else
  act sh -c "mkdir -p '$(dirname "$GRUB_CFG")' && printf '%s\n' 'GRUB_CMDLINE_LINUX_DEFAULT=\"\$GRUB_CMDLINE_LINUX_DEFAULT $CRASHKERNEL\"' > '$GRUB_CFG'"
fi
act update-grub || refuse "update-grub failed; not rebooting"
act systemctl enable kdump-tools.service || log "could not enable kdump-tools"

# one-shot verify unit, then reboot
if [ "$MODE" = dry-run ]; then
  log "would install $UNIT to verify kdump after boot, then reboot once"
  log "summary result=ok reboot=would kdump=pending"
  exit 0
fi
install -m 755 "$0" "$STATE/task.sh"
cat > "$UNIT" << UNITEOF
[Unit]
Description=Fillrate diag_tools_extended post-reboot kdump check
After=multi-user.target kdump-tools.service

[Service]
Type=oneshot
ExecStart=/bin/sh $STATE/task.sh verify

[Install]
WantedBy=multi-user.target
UNITEOF
act systemctl daemon-reload
act systemctl enable "$(basename "$UNIT")" || refuse "could not enable verify unit; not rebooting"
# Re-check rentals right before the reboot.
[ -z "$(docker ps -q 2>/dev/null)" ] || refuse "a container started during the run; not rebooting"
log "summary result=ok reboot=yes kdump=pending (verify unit writes $STATE/kdump-status after boot)"
act systemctl reboot
exit 0