#!/bin/sh
# fillrate-task: diag_access_setup
# One-time setup so Fillrate can run read-only host diagnostics without
# asking the host to run commands:
#   1. usermod -aG adm,systemd-journal fillrate  (read-only journal and
#      kernel log access; no sudo rule is added or changed)
#   2. install and enable rasdaemon; give the adm group read access to its
#      event DB (tmpfiles.d + rasdaemon drop-in, survives restarts and DB
#      recreation); load at most one EDAC driver, and only one the kernel's
#      modalias tables match to this hardware; unload unbound EDAC modules
#   3. install ipmitool only when an IPMI/BMC device is present
#   4. verify, as the fillrate user: journalctl -k and ras-mc-ctl --summary
#
# Safety:
# - Never restarts docker, vastai, nvidia-persistenced, or any NVIDIA/driver
#   service, and never reboots. needrestart is put in list-only mode and
#   policy-rc.d is not touched, so the only service apt starts is rasdaemon.
# - apt runs only after an apt-get -s simulation passes plan_check: new
#   installs, plus revision-only upgrades of dependencies (same epoch and
#   upstream version, e.g. libsqlite3-0 3.37.2-2ubuntu0.7 -> 0.8), each
#   logged as allow-upgrade. No removals, and nothing matching nvidia*,
#   libnvidia*, cuda*, docker*, containerd*, vast*, linux-image/modules/
#   headers*, libc6/libc-bin/locales, or systemd/udev. Otherwise the run
#   refuses and changes nothing.
# - Refuses when another apt/dpkg process holds the lock, or the ticket
#   expiry has passed.
# - Idempotent. A step that is already done is logged as already-done.
# - Ticket mode dry-run (FILLRATE_MODE=dry-run) runs every check and the
#   apt simulation and prints what it would do. It changes nothing.
#   An empty mode or apply makes the changes.
# - Every step is logged to stdout and to a mode 644 log under
#   /var/lib/fillrate-runner/diag-access-setup. No keys, tokens, or ticket
#   nonce are written.
# The runner starts it with a clean environment. The FILLRATE_* path
# overrides below are for tests; a real run uses the defaults.

set -u

STATE=${FILLRATE_DIAG_STATE:-/var/lib/fillrate-runner/diag-access-setup}
MODULES_LOAD=${FILLRATE_MODULES_LOAD:-/etc/modules-load.d/fillrate-edac.conf}
EDAC_SYS=${FILLRATE_EDAC_SYS:-/sys/devices/system/edac/mc}
SYSFS=${FILLRATE_SYSFS:-/sys}
PROC_MODULES=${FILLRATE_PROC_MODULES:-/proc/modules}
RAS_TMPFILES=${FILLRATE_RAS_TMPFILES:-/etc/tmpfiles.d/fillrate-rasdaemon.conf}
RAS_DROPIN=${FILLRATE_RAS_DROPIN:-/etc/systemd/system/rasdaemon.service.d/fillrate-perms.conf}
DEV_DIR=${FILLRATE_DEV_DIR:-/dev}
EXPIRES=${FILLRATE_EXPIRES:-0}
MODE=${FILLRATE_MODE:-apply}
RUNNER_USER=fillrate

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

case "$MODE" in
  ""|apply) MODE=apply ;;
  dry-run) ;;
  *) printf '%s\n' "refusing: mode must be apply or dry-run"; exit 1 ;;
esac

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

GROUPS_RESULT=unknown
RAS_RESULT=unknown
EDAC_RESULT=unknown
IPMI_RESULT=unknown
VERIFY_JOURNAL=unknown
VERIFY_RAS=unknown
FAILED=0

log() {
  line="$(date -u +%H:%M:%SZ 2>/dev/null) [$MODE] $*"
  printf '%s\n' "$line"
  printf '%s\n' "$line" >> "$LOG" 2>/dev/null || true
}

refuse() {
  log "refusing: $1"
  summary refused
  exit 1
}

summary() {
  log "summary result=$1 groups=$GROUPS_RESULT rasdaemon=$RAS_RESULT ras_db_access=$RAS_PERMS edac=$EDAC_RESULT ipmitool=$IPMI_RESULT verify_journal=$VERIFY_JOURNAL verify_ras=$VERIFY_RAS log=$LOG"
}

run() {
  # Log the command, then run it unless this is a dry run.
  if [ "$MODE" = dry-run ]; then
    log "would run: $*"
    return 0
  fi
  log "run: $*"
  out=$("$@" 2>&1)
  code=$?
  [ -n "$out" ] && printf '%s\n' "$out" | sed 's/^/    /' | tee -a "$LOG" 2>/dev/null
  log "exit=$code: $1"
  return $code
}

window_open() {
  [ "$EXPIRES" = 0 ] && return 0
  now=$(date +%s) || return 1
  [ "$now" -le "$EXPIRES" ]
}

pkg_installed() {
  dpkg-query -W -f='${Status}' "$1" 2>/dev/null | grep -q 'install ok installed'
}

ipmi_present() {
  for d in "$DEV_DIR"/ipmi0 "$DEV_DIR"/ipmi/0 "$DEV_DIR"/ipmidev/0; do
    [ -e "$d" ] && return 0
  done
  if command -v dmidecode >/dev/null 2>&1; then
    dmidecode -t 38 2>/dev/null | grep -q 'IPMI Device Information' && return 0
  fi
  return 1
}

# 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 }
  '
}

# Simulate an apt install and run plan_check on it. Returns non-zero if unsafe.
safe_plan() {
  plan=$(apt-get -s -o Debug::NoLocking=1 install --no-install-recommends "$@" 2>&1)
  code=$?
  if [ $code -ne 0 ]; then
    log "apt simulation failed (exit $code)"
    printf '%s\n' "$plan" | tail -n 20 | sed 's/^/    /' | tee -a "$LOG" 2>/dev/null
    return 1
  fi
  verdict=$(printf '%s\n' "$plan" | plan_check "$@")
  vcode=$?
  printf '%s\n' "$verdict" | sed '/^$/d; s/^/    plan: /' | tee -a "$LOG" 2>/dev/null
  return $vcode
}

edac_say() { log "$*"; }
edac_act() { run "$@"; }
EDAC_MODE=$MODE
EDAC_UNLOADED=""
RAS_PERMS=skipped

# ---- 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
}

log "diag_access_setup start machine=${FILLRATE_MACHINE_ID:-unknown} request=${FILLRATE_BM_REQUEST_ID:-unknown}"

# ---- prechecks ----
[ "$(id -u)" -eq 0 ] || refuse "not root"
window_open || refuse "ticket expiry has passed"
id "$RUNNER_USER" >/dev/null 2>&1 || refuse "user $RUNNER_USER does not exist"
command -v apt-get >/dev/null 2>&1 || refuse "apt-get not found (only Debian/Ubuntu hosts are supported)"
if command -v fuser >/dev/null 2>&1; then
  if fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock >/dev/null 2>&1; then
    refuse "another apt/dpkg process holds the lock"
  fi
fi
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
running=$(docker ps -q 2>/dev/null | wc -l | tr -d ' ')
log "precheck ok: running containers=${running:-unknown} (left untouched; nothing restarts docker, vastai, or the driver)"

# Decide what to install before touching anything.
want=""
pkg_installed rasdaemon || want="rasdaemon"
if ipmi_present; then
  log "ipmi device present"
  pkg_installed ipmitool && IPMI_RESULT=already-installed || want="$want ipmitool"
else
  IPMI_RESULT=skipped-no-bmc
  log "no IPMI/BMC device; ipmitool skipped"
fi
want=$(printf '%s' "$want" | sed 's/^ *//')

if [ -n "$want" ]; then
  if ! apt-cache policy $want 2>/dev/null | grep -q 'Candidate: [0-9]'; then
    run apt-get update -q || refuse "apt-get update failed"
  fi
  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"
fi

# ---- 1. groups ----
missing=""
for g in adm systemd-journal; do
  getent group "$g" >/dev/null 2>&1 || { log "group $g does not exist; skipped"; continue; }
  id -nG "$RUNNER_USER" | tr ' ' '\n' | grep -x -q "$g" || missing="$missing,$g"
done
missing=${missing#,}
if [ -z "$missing" ]; then
  GROUPS_RESULT=already-done
  log "groups: $RUNNER_USER already in adm and systemd-journal"
elif run usermod -aG "$missing" "$RUNNER_USER"; then
  GROUPS_RESULT=added:$missing
else
  GROUPS_RESULT=failed; FAILED=1
fi

# ---- 2. packages (rasdaemon, ipmitool) ----
if [ -n "$want" ]; then
  if run apt-get install -y -q --no-install-recommends -o Dpkg::Options::=--force-confold $want; then
    case " $want " in *" rasdaemon "*) RAS_RESULT=installed ;; esac
    case " $want " in *" ipmitool "*) IPMI_RESULT=installed ;; esac
  else
    FAILED=1
    case " $want " in *" rasdaemon "*) RAS_RESULT=install-failed ;; esac
    case " $want " in *" ipmitool "*) IPMI_RESULT=install-failed ;; esac
  fi
fi
[ "$RAS_RESULT" = unknown ] && RAS_RESULT=already-installed
if [ "$MODE" = dry-run ] || pkg_installed rasdaemon; then
  if systemctl is-enabled rasdaemon >/dev/null 2>&1 && systemctl is-active rasdaemon >/dev/null 2>&1; then
    log "rasdaemon already enabled and active"
  else
    run systemctl enable --now rasdaemon.service || { FAILED=1; RAS_RESULT="$RAS_RESULT,enable-failed"; }
  fi
fi

# ---- 3. rasdaemon DB access + EDAC ----
if [ "$MODE" = dry-run ] || pkg_installed rasdaemon; then
  ras_perms_setup || { RAS_PERMS=failed; FAILED=1; }
  log "rasdaemon DB access for adm: $RAS_PERMS"
fi
edac_setup
[ -n "$EDAC_UNLOADED" ] && log "EDAC modules unloaded:$EDAC_UNLOADED"

# ---- 4. verify as fillrate ----
if [ "$MODE" = dry-run ]; then
  log "would verify: runuser -u $RUNNER_USER -- journalctl -k -n 5 --no-pager; runuser -u $RUNNER_USER -- ras-mc-ctl --summary"
  VERIFY_JOURNAL=skipped-dry-run
  VERIFY_RAS=skipped-dry-run
else
  jout=$(runuser -u "$RUNNER_USER" -- journalctl -k -n 5 --no-pager -q 2>&1)
  jcode=$?
  printf '%s\n' "$jout" | sed 's/^/    /' >> "$LOG" 2>/dev/null
  if [ $jcode -eq 0 ] && ! printf '%s\n' "$jout" | grep -q -i 'not seeing messages\|permission denied'; then
    VERIFY_JOURNAL=ok
  else
    VERIFY_JOURNAL=failed; FAILED=1
  fi
  log "verify journalctl -k as $RUNNER_USER: $VERIFY_JOURNAL"
  rout=$(runuser -u "$RUNNER_USER" -- ras-mc-ctl --summary 2>&1)
  rcode=$?
  printf '%s\n' "$rout" | sed 's/^/    /' | tee -a "$LOG" 2>/dev/null
  if [ $rcode -eq 0 ] && ! printf '%s\n' "$rout" | grep -q -i 'unable to open\|failed\|can.t call'; then
    VERIFY_RAS=ok
  else
    VERIFY_RAS=failed; FAILED=1
  fi
  log "verify ras-mc-ctl --summary as $RUNNER_USER: $VERIFY_RAS"
  log "note: SSH sessions opened before this run keep the old groups; reconnect to pick them up"
fi

if [ "$FAILED" -ne 0 ]; then
  summary partial
  exit 2
fi
summary ok
exit 0