#!/usr/bin/env bash
# Install / refresh the CIS490 lab-host role.
#
# Idempotent — safe to re-run after `git pull`. Does NOT enroll the
# host into WireGuard (that's wg-enroll's job, run separately and
# *first*) and does NOT mint TLS certs (that's wg-pki's job).
#
# Steps:
#   1. Verify prereqs (KVM, zstd, qemu, python3.11+, systemd).
#   2. Create the cis490 service user + /var/lib/cis490 layout.
#   3. Sync the repo into /opt/cis490 and build a uv-managed venv.
#   4. Install systemd units from etc/.
#   5. Drop /etc/cis490/lab-host.toml (only on first install).
#
# Operator finishes by:
#   - editing /etc/cis490/lab-host.toml (host_id, receiver URL, certs)
#   - placing leaf certs at /etc/cis490/certs/{lab-host.pem,key,wg-ca.pem}
#   - `systemctl enable --now cis490-shipper`

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
INSTALL_ROOT="${INSTALL_ROOT:-/opt/cis490}"
DATA_ROOT="${DATA_ROOT:-/var/lib/cis490}"
ETC_ROOT="${ETC_ROOT:-/etc/cis490}"
SERVICE_USER="${SERVICE_USER:-cis490}"

log() { printf '[install-lab-host] %s\n' "$*" >&2; }
die() { log "FATAL: $*"; exit 1; }

# Detect distro family from /etc/os-release. Returns one of
# "arch" | "debian" | "rhel" | "other". Used to install collector
# binaries (perf, tcpdump) that the collectors need but which aren't
# always pre-installed. Per PIPELINE.md §1, a host that can't host the
# canonical experiment must produce zero episodes — this function is
# the install-time path to actually installing the deps so the host
# CAN host the experiment, rather than the bandaid path of running
# without them.
detect_os() {
    [[ -f /etc/os-release ]] || { echo other; return; }
    # shellcheck disable=SC1091
    . /etc/os-release
    case "${ID:-}${ID_LIKE:-}" in
        *arch*)             echo arch ;;
        *debian*|*ubuntu*)  echo debian ;;
        *rhel*|*fedora*|*centos*) echo rhel ;;
        *)                  echo other ;;
    esac
}

# Install collector binaries that aren't already on PATH. Idempotent;
# `--needed` / equivalent skip already-installed packages. We do NOT
# do a full system upgrade (-Syu / dist-upgrade) — operator owns that.
# If a package install fails the script falls through to the existing
# `command -v` checks below, which die loudly with a clear message.
ensure_collector_packages() {
    local need_perf=0 need_tcpdump=0
    command -v perf    >/dev/null || need_perf=1
    command -v tcpdump >/dev/null || need_tcpdump=1
    [[ $need_perf -eq 0 && $need_tcpdump -eq 0 ]] && return 0

    local os; os=$(detect_os)
    log "missing collector binaries (perf=$need_perf tcpdump=$need_tcpdump); os=$os"
    case "$os" in
        arch)
            # Arch ships perf in the `perf` package (from extra). The
            # canonical kernel is `linux`; if the operator runs a
            # different kernel flavor (linux-lts, linux-zen) the perf
            # package is the same — perf is kernel-version-aware via
            # `uname -r` at runtime.
            local pkgs=()
            [[ $need_perf    -eq 1 ]] && pkgs+=(perf)
            [[ $need_tcpdump -eq 1 ]] && pkgs+=(tcpdump)
            log "pacman -S --needed --noconfirm ${pkgs[*]}"
            pacman -S --needed --noconfirm "${pkgs[@]}" \
                || log "WARN: pacman install failed; falling through to die-on-missing"
            ;;
        debian)
            local pkgs=()
            # On Debian/Ubuntu perf comes from linux-perf or
            # linux-tools-generic depending on release; try both.
            [[ $need_perf    -eq 1 ]] && pkgs+=(linux-perf linux-tools-generic)
            [[ $need_tcpdump -eq 1 ]] && pkgs+=(tcpdump)
            log "apt-get install -y ${pkgs[*]}"
            DEBIAN_FRONTEND=noninteractive apt-get install -y "${pkgs[@]}" \
                || log "WARN: apt-get install failed; falling through to die-on-missing"
            ;;
        rhel)
            local pkgs=()
            [[ $need_perf    -eq 1 ]] && pkgs+=(perf)
            [[ $need_tcpdump -eq 1 ]] && pkgs+=(tcpdump)
            log "dnf install -y ${pkgs[*]}"
            dnf install -y "${pkgs[@]}" \
                || log "WARN: dnf install failed; falling through to die-on-missing"
            ;;
        *)
            log "WARN: unknown distro; cannot auto-install perf/tcpdump"
            log "      install them by hand for collectors §4.4 to pass"
            ;;
    esac
}

# --- 1. prereqs --------------------------------------------------------
log "checking prereqs"

if [[ $EUID -ne 0 ]]; then
    die "must run as root (writes to /opt, /etc, /var/lib, and systemd)"
fi
command -v systemctl >/dev/null || die "systemd not found"
command -v qemu-system-x86_64 >/dev/null || die "qemu-system-x86_64 not on PATH"
command -v zstd >/dev/null || die "zstd not on PATH (install via your package manager)"
[[ -e /dev/kvm ]] || die "/dev/kvm missing — KVM not available"

# Auto-install collector binaries that aren't on PATH. Done after the
# always-required checks above (qemu, kvm, zstd, systemd) which the
# operator has to provide themselves at OS install time.
ensure_collector_packages

# Re-check perf/tcpdump after the install attempt and die loudly if
# still missing. Per §4.4 collectors that can't run shouldn't ship
# silently — the host fails install instead so operator notices.
command -v perf >/dev/null || die \
    "perf not on PATH after install attempt — collector source 3 (oracle perf-stat) requires it. See ensure_collector_packages above for what was tried."
command -v tcpdump >/dev/null || die \
    "tcpdump not on PATH after install attempt — collector source 4 (bridge pcap + netflow) requires it. See ensure_collector_packages above for what was tried."

# Canonical experiment manifest (PIPELINE.md §4.1). Validate at install
# time so a host that cloned a repo missing manifest.toml fails the
# install loudly, not at first orchestrator startup. The orchestrator
# also validates and exits 78 on bad manifest, but install-time is the
# earliest point we can fail.
[[ -f "$REPO_ROOT/manifest.toml" ]] || die \
    "$REPO_ROOT/manifest.toml not found — every lab host must ship the canonical experiment manifest (§4.1). Did the git pull complete?"
"$REPO_ROOT/.venv/bin/python" -c "
import sys, pathlib
sys.path.insert(0, '$REPO_ROOT')
from orchestrator.manifest import load_canonical, ManifestError
try:
    load_canonical(pathlib.Path('$REPO_ROOT'))
except ManifestError as e:
    print(f'manifest.toml failed validation: {e}', file=sys.stderr)
    sys.exit(1)
" || die "manifest.toml validation failed — see error above (§4.1)"

# uv is preferred (lockfile-driven). Fall back to system pip if absent.
USE_UV=0
if command -v uv >/dev/null; then USE_UV=1; fi

# --- 2. user + layout --------------------------------------------------
log "ensuring service user $SERVICE_USER"
if ! id -u "$SERVICE_USER" >/dev/null 2>&1; then
    useradd --system --no-create-home --shell /usr/sbin/nologin \
        --home-dir "$INSTALL_ROOT" "$SERVICE_USER"
fi
# kvm group lets the service spawn VMs.
if getent group kvm >/dev/null 2>&1; then
    usermod -a -G kvm "$SERVICE_USER" || true
fi

install -d -o root -g root -m 0755 "$ETC_ROOT" "$ETC_ROOT/certs"
install -d -o "$SERVICE_USER" -g "$SERVICE_USER" -m 0755 \
    "$DATA_ROOT" "$DATA_ROOT/data" \
    "$DATA_ROOT/data/episodes" "$DATA_ROOT/data/outbox" \
    "$DATA_ROOT/data/shipped" "$DATA_ROOT/data/queue" \
    "$DATA_ROOT/samples" "$DATA_ROOT/samples/store" \
    "$DATA_ROOT/vm" "$DATA_ROOT/vm/images"

# --- 3. repo + venv ----------------------------------------------------
log "syncing repo into $INSTALL_ROOT"
install -d -o "$SERVICE_USER" -g "$SERVICE_USER" -m 0755 "$INSTALL_ROOT"
# The expected lab-host upgrade path is `cd /opt/cis490 && git pull &&
# sudo ./scripts/install-lab-host.sh`, which means REPO_ROOT and
# INSTALL_ROOT can be the same directory. cp -aT errors out in that
# case ("are the same file") and `set -e` aborts before the systemd
# units get installed and services restart, leaving the host running
# whatever code was loaded at last service start.
if [[ "$(readlink -f "$REPO_ROOT")" == "$(readlink -f "$INSTALL_ROOT")" ]]; then
    log "REPO_ROOT == INSTALL_ROOT ($INSTALL_ROOT); git pull already updated tree, skipping cp"
else
    # Clean cp -aT to avoid an extra rsync dep.
    cp -aT "$REPO_ROOT" "$INSTALL_ROOT"
fi
chown -R "$SERVICE_USER":"$SERVICE_USER" "$INSTALL_ROOT"

# Stamp a VERSION file at install time so episodes can record the
# code commit they were generated by. /opt/cis490 is a flat copy
# (no .git/), so we capture the source repo's HEAD here. Trainers
# read meta.json.code_version to filter out episodes from buggy
# pre-fix code.
if VC="$(cd "$REPO_ROOT" && git rev-parse HEAD 2>/dev/null)"; then
    VB="$(cd "$REPO_ROOT" && git rev-parse --abbrev-ref HEAD 2>/dev/null || echo unknown)"
    VD="false"
    if cd "$REPO_ROOT" && [[ -n "$(git status --porcelain 2>/dev/null)" ]]; then
        VD="true"
    fi
    install -o "$SERVICE_USER" -g "$SERVICE_USER" -m 0644 /dev/stdin \
        "$INSTALL_ROOT/VERSION" <<EOF
{"commit": "$VC", "branch": "$VB", "dirty": $VD, "installed_at_wall": "$(date -u +%Y-%m-%dT%H:%M:%SZ)"}
EOF
    log "VERSION stamp: $VC ($VB)$([[ "$VD" == "true" ]] && echo " [dirty]")"
else
    log "WARN: $REPO_ROOT not a git checkout; episodes will record code_version.commit='unknown'"
fi

log "building venv"
if [[ "$USE_UV" -eq 1 ]]; then
    sudo -u "$SERVICE_USER" -- env HOME="$INSTALL_ROOT" \
        uv sync --project "$INSTALL_ROOT"
else
    sudo -u "$SERVICE_USER" -- python3 -m venv "$INSTALL_ROOT/.venv"
    sudo -u "$SERVICE_USER" -- "$INSTALL_ROOT/.venv/bin/pip" install \
        --quiet --upgrade pip
    sudo -u "$SERVICE_USER" -- "$INSTALL_ROOT/.venv/bin/pip" install \
        --quiet starlette 'uvicorn[standard]' httpx msgpack
fi

# --- 4. systemd --------------------------------------------------------
log "installing systemd units"
install -m 0644 "$REPO_ROOT/etc/cis490-shipper.service" \
    /etc/systemd/system/cis490-shipper.service
install -m 0644 "$REPO_ROOT/etc/cis490-orchestrator.service" \
    /etc/systemd/system/cis490-orchestrator.service
# Auto-update: a 30-min timer that does git fetch + (if behind) pull
# and re-run this script. Prevents host-falls-behind incidents when
# the receiver's allow-list rolls forward and an on-device agent
# fails to act on the 412 remediation. See AGENTS.md "Auto-update".
install -m 0644 "$REPO_ROOT/etc/cis490-autoupdate.service" \
    /etc/systemd/system/cis490-autoupdate.service
install -m 0644 "$REPO_ROOT/etc/cis490-autoupdate.timer" \
    /etc/systemd/system/cis490-autoupdate.timer
# mTLS cert-fetch retry timer: handles the "operator edited host_id
# but didn't re-run install" case AND the "bootstrap.wg was briefly
# unreachable" case. Polls every 5 min; no-op once the cert is on
# disk. See scripts/fetch-lab-host-cert.sh.
install -m 0644 "$REPO_ROOT/etc/cis490-cert-fetch.service" \
    /etc/systemd/system/cis490-cert-fetch.service
install -m 0644 "$REPO_ROOT/etc/cis490-cert-fetch.timer" \
    /etc/systemd/system/cis490-cert-fetch.timer
# Daily doctor check — runs cis490_doctor.py and PUTs JSON to the
# receiver's /v1/host-health/<host>. Lets the maintainer see fleet
# health without grepping per-host logs. See AGENTS.md "Fleet health".
install -m 0644 "$REPO_ROOT/etc/cis490-doctor-check.service" \
    /etc/systemd/system/cis490-doctor-check.service
install -m 0644 "$REPO_ROOT/etc/cis490-doctor-check.timer" \
    /etc/systemd/system/cis490-doctor-check.timer
systemctl daemon-reload
# Enable timers immediately — the operator gets self-healing on the
# next tick without an extra `systemctl enable`. Idempotent.
for t in cis490-autoupdate.timer cis490-cert-fetch.timer cis490-doctor-check.timer; do
    systemctl enable --now "$t" 2>/dev/null || \
        log "WARN: could not enable $t (will retry next install)"
done

# --- 5. config template (only on first install) -----------------------
if [[ ! -f "$ETC_ROOT/lab-host.toml" ]]; then
    log "writing $ETC_ROOT/lab-host.toml (template)"
    install -m 0640 -o root -g "$SERVICE_USER" \
        "$REPO_ROOT/etc/lab-host.toml.example" "$ETC_ROOT/lab-host.toml"
    NEW_INSTALL=1
else
    log "$ETC_ROOT/lab-host.toml exists; leaving in place"
    NEW_INSTALL=0
fi

# --- 6. orchestrator env file (read by cis490-orchestrator.service) ----
ENV_FILE="$ETC_ROOT/lab-host.env"
DEFAULT_HOST_ID="$(hostname -s)"
if [[ ! -f "$ENV_FILE" ]]; then
    log "writing $ENV_FILE (host_id defaults to $DEFAULT_HOST_ID — edit if you want something else)"
    install -m 0640 -o root -g "$SERVICE_USER" /dev/stdin "$ENV_FILE" <<EOF
# Read by cis490-orchestrator.service. Override per-host as needed.
FLEET_HOST_ID=$DEFAULT_HOST_ID
# BRIDGE=br-malware enables source 4 pcap capture AND unlocks the
# Tier-3 modules whose payloads need callback (reverse/bind shells).
# install-lab-host.sh provisions the bridge + tap pool below; leave
# this on unless your lab host can't run NETLINK ops.
BRIDGE=br-malware
EOF
fi

# --- 6b. host-only bridge + per-slot tap pool --------------------------
# br-malware lets pcap capture the guest traffic and lets bind/reverse
# shell payloads route between guest and host. We pre-create a small
# pool of taps so the launchers don't need sudo to attach interfaces;
# each slot uses cis490tap{SLOT,SLOT+100} (Tier-2 demo + Tier-3
# target). Idempotent: re-running on an already-set-up host is a
# no-op.
if command -v ip >/dev/null && [[ -x "$REPO_ROOT/vm/setup_bridge.sh" ]]; then
    if "$REPO_ROOT/vm/setup_bridge.sh" >/dev/null 2>&1; then
        log "bridge br-malware ready"
        for n in 0 1 2 3 4 5 6 7; do
            for prefix in cis490tap cis490target; do
                tap="${prefix}${n}"
                if ! ip link show "$tap" >/dev/null 2>&1; then
                    ip tuntap add dev "$tap" mode tap user "$SERVICE_USER" 2>/dev/null || \
                        ip tuntap add dev "$tap" mode tap 2>/dev/null || true
                    ip link set "$tap" master br-malware 2>/dev/null || true
                    ip link set "$tap" up 2>/dev/null || true
                fi
            done
        done
        log "tap pool: cis490tap0..7 + cis490target0..7 attached to br-malware"
    else
        log "WARN: setup_bridge.sh failed; BRIDGE mode will be unavailable"
        # Comment out BRIDGE in the env file — fleet will still run
        # Tier-2 + non-callback Tier-3 modules.
        sed -i 's/^BRIDGE=br-malware/# BRIDGE=br-malware  # auto-disabled: bridge setup failed/' "$ENV_FILE"
    fi
fi

# --- 7. mTLS leaf cert (auto-fetch via bootstrap.wg) -------------------
# One-shot fetch via the standalone script (also wired to a 5-min
# retry timer in step 4 above, so this is just the install-time
# kick-off — the timer handles transient failures and the case
# where the operator edits host_id later).
"$INSTALL_ROOT/scripts/fetch-lab-host-cert.sh" || \
    log "WARN: cert fetch failed — timer will retry"

# --- 8. baseline VM image + cidata (best-effort) -----------------------
ALPINE_IMG="$DATA_ROOT/vm/images/alpine-baseline.qcow2"
CIDATA_ISO="$DATA_ROOT/vm/images/cidata.iso"
if [[ ! -f "$ALPINE_IMG" ]]; then
    if "$REPO_ROOT/scripts/fetch-alpine-baseline.sh" "$ALPINE_IMG"; then
        log "fetched Alpine baseline -> $ALPINE_IMG"
    else
        log "WARN: Alpine baseline fetch failed; drop a qcow2 at $ALPINE_IMG manually"
    fi
fi
if [[ -f "$ALPINE_IMG" && ! -f "$CIDATA_ISO" ]]; then
    log "building cidata.iso (in-guest agent embedded)"
    sudo -u "$SERVICE_USER" -- "$INSTALL_ROOT/.venv/bin/python" \
        "$INSTALL_ROOT/tools/build_cidata.py" "$CIDATA_ISO" || \
        log "WARN: cidata build failed; run tools/build_cidata.py manually"
fi
# Symlink the canonical paths the launchers look at, when missing.
install -d -o "$SERVICE_USER" -g "$SERVICE_USER" -m 0755 "$INSTALL_ROOT/vm/images"
ln -sf "$ALPINE_IMG" "$INSTALL_ROOT/vm/images/alpine-baseline.qcow2" 2>/dev/null || true
ln -sf "$CIDATA_ISO" "$INSTALL_ROOT/vm/images/cidata.iso" 2>/dev/null || true
M2_IMG="$DATA_ROOT/vm/images/metasploitable2.qcow2"
[[ -f "$M2_IMG" ]] && ln -sf "$M2_IMG" "$INSTALL_ROOT/vm/images/metasploitable2.qcow2" 2>/dev/null || true

# --- 8. Tier-3 + Tier-4 deploy (auto, idempotent) ----------------------
# Bring up msfrpcd + Metasploitable2 + bridge + verify. Skipped only if
# certs aren't on disk yet (Tier-3 fire writes episodes that the
# shipper ships, so it's pointless to run before mTLS is live) or the
# operator passed --skip-tier3.
SKIP_TIER3="${SKIP_TIER3:-}"
for arg in "$@"; do
    [[ "$arg" == "--skip-tier3" ]] && SKIP_TIER3=1
done
if [[ -z "$SKIP_TIER3" && -f "$ETC_ROOT/certs/lab-host.pem" ]]; then
    log "deploying Tier 3 (msfrpcd + Metasploitable2 + bridge)"
    if "$INSTALL_ROOT/scripts/install-tier-3-4.sh"; then
        log "Tier-3 deploy ✓"
    else
        log "WARN: Tier-3 deploy failed — Tier 2 will keep running."
        log "      Re-run later: sudo $INSTALL_ROOT/scripts/install-tier-3-4.sh"
    fi
elif [[ -z "$SKIP_TIER3" ]]; then
    log "skipping Tier-3 deploy (no mTLS cert yet — re-run this script after"
    log "host_id is set so the cert auto-fetches first)"
fi

if [[ "$NEW_INSTALL" == "1" ]]; then
    log ""
    log "================================================================="
    log "  FIRST-INSTALL NEXT STEPS                                        "
    log "================================================================="
    log "  1. Edit $ETC_ROOT/lab-host.toml — set host_id and receiver URL."
    log "     (host_id starts as 'REPLACE_ME'; the cert auto-fetch in"
    log "     step 7 of this script SKIPS while that's still the value.)"
    log ""
    log "  2. RE-RUN THIS SCRIPT — sudo $0"
    log "     The second pass sees the real host_id and pulls the leaf"
    log "     cert from https://bootstrap.wg/v1/cert/<host_id>. There is"
    log "     no manual cert-minting step on this host. DO NOT openssl."
    log ""
    log "  3. Confirm certs landed:"
    log "       ls -l $ETC_ROOT/certs/"
    log "     Expected: lab-host.pem, lab-host.key, wg-ca.pem"
    log ""
    log "  4. Run the diagnostic — every red row prints the exact fix:"
    log "       $INSTALL_ROOT/.venv/bin/python \\"
    log "           $INSTALL_ROOT/tools/cis490_doctor.py --role lab-host"
    log ""
    log "  5. Smoke-test the pipe (returns ok=true on success):"
    log "       sudo -u $SERVICE_USER $INSTALL_ROOT/.venv/bin/python -m shipper \\"
    log "            --config $ETC_ROOT/lab-host.toml --ping"
    log ""
    log "  6. Turn on the services — episodes start flowing immediately:"
    log "       sudo systemctl enable --now cis490-shipper cis490-orchestrator"
    log ""
    log "  If bootstrap.wg fetch fails (step 2 prints WARN), see AGENTS.md"
    log "  ('Securing the connection (mTLS)') — almost always a missing"
    log "  /etc/hosts entry or wg0 not up. Do not work around this by"
    log "  minting certs locally; fix the WG path instead."
    log "================================================================="
fi

# --- 9. drain pre-stamp queue + restart services -----------------------
# On a re-run (not first install), the running services are still
# executing whatever code they loaded at last start, so the new module
# objects we just dropped on disk don't take effect until restart. And
# any episodes the orchestrator generated against pre-stamping code
# are missing meta.json::code_version — the receiver returns 400 for
# every PUT, the queue's fatal path didn't quarantine them in older
# code, and the shipper burns cycles re-tarring them on every pass.
# Drain them once, then restart so the new code reaches the live
# daemon.
if [[ "$NEW_INSTALL" != "1" ]]; then
    PY="$INSTALL_ROOT/.venv/bin/python"
    if [[ -x "$PY" && -f "$INSTALL_ROOT/tools/quarantine_unstamped.py" ]]; then
        log "draining pre-stamp episodes from queue (idempotent)"
        sudo -u "$SERVICE_USER" -- "$PY" \
            "$INSTALL_ROOT/tools/quarantine_unstamped.py" \
            --data-root "$DATA_ROOT/data" || \
            log "WARN: quarantine drain returned non-zero — see output above"
    fi
    systemctl daemon-reload
    for svc in cis490-shipper cis490-orchestrator; do
        if systemctl is-enabled --quiet "$svc" 2>/dev/null; then
            log "restarting $svc to pick up new code"
            systemctl restart "$svc" || \
                log "WARN: $svc restart failed — check 'journalctl -u $svc'"
        fi
    done
fi

log "lab-host install complete."
log ""
log "Cloning this repo and running the launchers manually is NOT enough."
log "The lab-host role's data flow lives in the systemd services this"
log "script just installed. If $INSTALL_ROOT/index.jsonl on the Pi stays"
log "empty after step 5, run:"
log "   $INSTALL_ROOT/.venv/bin/python $INSTALL_ROOT/tools/cis490_doctor.py"