#!/usr/bin/env bash
# Tests for bin/_agent-unsafe.sh and its cld-unsafe / cdx-unsafe / ca-unsafe symlinks.
#
# The allowed path is exercised with a STUB runtime under a fake HOME and a real pty
# (script -qec), with every agent env marker stripped via `env -u`. There is deliberately
# no test-mode flag in the production script: any such flag would be settable by the very
# agent the gate exists to refuse.
#
# Run: bash agent-unsafe.test.sh (exit 0 = all pass).
set -uo pipefail
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
LIB="$ROOT/bin/_agent-unsafe.sh"
SESSION_LIB="$ROOT/lib/session-class.sh"
PASS=0; FAIL=0
ok()  { PASS=$((PASS+1)); printf 'PASS %s\n' "$1"; }
bad() { FAIL=$((FAIL+1)); printf 'FAIL %s\n     %s\n' "$1" "$2"; }

TMP=$(mktemp -d "${TMPDIR:-/tmp}/agent-unsafe-test-XXXX")
trap 'rm -rf "$TMP"' EXIT

# Marker list comes from the production script, so a marker added later cannot rot the test.
# shellcheck source=/dev/null
source "$LIB"
UNSET_ARGS=()
for m in "${SESSION_AGENT_MARKERS[@]}"; do UNSET_ARGS+=(-u "$m"); done

FAKE_HOME="$TMP/home"
mkdir -p "$FAKE_HOME/.claude/bin"
for r in claude codex; do
  cat >"$FAKE_HOME/.claude/bin/$r" <<'STUB'
#!/usr/bin/env bash
for a in "$@"; do printf '%s\0' "$a" >>"$STUB_ARGV"; done
awk -F: '$1=="0"{print $3}' /proc/self/cgroup >"$STUB_CGROUP"
[[ -n "${STUB_SLEEP:-}" ]] && sleep "$STUB_SLEEP"
exit 0
STUB
  chmod +x "$FAKE_HOME/.claude/bin/$r"
done

# Runs a hatch on a real pty with the agent markers stripped. $1=hatch, rest=argv.
# Prints combined output; returns the hatch's exit code.
run_on_tty() {
  local hatch="$1"; shift
  local runner="$TMP/runner.$$.$RANDOM.sh" a
  {
    printf '#!/usr/bin/env bash\nexec env'
    printf ' %s' "${UNSET_ARGS[@]}"
    printf ' HOME=%q %q' "$HOME_FOR_RUN" "$ROOT/bin/$hatch"
    for a in "$@"; do printf ' %q' "$a"; done
    printf '\n'
  } >"$runner"
  chmod +x "$runner"
  script -qec "$runner" /dev/null
}

LOG_REL=".local/state/agent-unsafe/invocations.log"

# (a) refuses when AGENT_BUILD_SCOPE_ACTIVE is set
HOME_FOR_RUN="$FAKE_HOME"
out=$(env HOME="$FAKE_HOME" AGENT_BUILD_SCOPE_ACTIVE=1 "$ROOT/bin/cld-unsafe" x 2>&1); rc=$?
if (( rc == 77 )) && [[ "$out" == *"agent-env-marker:AGENT_BUILD_SCOPE_ACTIVE"* ]]; then
  ok "a: refuses when AGENT_BUILD_SCOPE_ACTIVE is set"
else
  bad "a: refuses when AGENT_BUILD_SCOPE_ACTIVE is set" "rc=$rc out=$out"
fi

# (b) refuses when stdin is not a TTY (markers stripped, so no-tty is the only trip)
out=$(env "${UNSET_ARGS[@]}" HOME="$FAKE_HOME" "$ROOT/bin/cld-unsafe" x </dev/null 2>&1); rc=$?
if (( rc == 77 )) && [[ "$out" == *"REFUSED (no-tty)"* ]]; then
  ok "b: refuses when stdin is not a TTY"
else
  bad "b: refuses when stdin is not a TTY" "rc=$rc out=$out"
fi

# (c1) refuses on the container= env signal, even at a real TTY
out=$(script -qec "env $(printf '%s ' "${UNSET_ARGS[@]}")container=podman HOME=$FAKE_HOME $ROOT/bin/cld-unsafe x" /dev/null 2>&1); rc=$?
if (( rc == 77 )) && [[ "$out" == *"REFUSED (in-container)"* && "$out" == *"env:container=podman"* ]]; then
  ok "c1: refuses on container= env marker"
else
  bad "c1: refuses on container= env marker" "rc=$rc out=$out"
fi

# (c2) refuses on a file/cgroup container signal (function overridden, not an env knob)
out=$(env "${UNSET_ARGS[@]}" HOME="$FAKE_HOME" bash -c '
  source "$1"
  session_container_signal() { printf "file:/run/.containerenv"; }
  unsafe_main cld-unsafe x
' _ "$LIB" 2>&1); rc=$?
if (( rc == 77 )) && [[ "$out" == *"REFUSED (in-container)"* && "$out" == *"/run/.containerenv"* ]]; then
  ok "c2: refuses on a file container signal"
else
  bad "c2: refuses on a file container signal" "rc=$rc out=$out"
fi
for sig in '/run/.containerenv' '/.dockerenv' 'container' '/proc/self/cgroup'; do
  grep -qF "$sig" "$SESSION_LIB" || bad "c2: container signal $sig referenced in $SESSION_LIB" "missing"
done

# (d) ca-unsafe still rejects a -fast cursor model id (real ca.sh, all required flags valid)
HOME_FOR_RUN="$HOME"
out=$(run_on_tty ca-unsafe --workspace "$TMP" --trust 'probe prompt' --task-slug fastprobe --model composer-1-fast 2>&1); rc=$?
if (( rc == 2 )) && [[ "$out" == *"REFUSED"* && "$out" == *"-fast tier"* ]]; then
  ok "d: ca-unsafe still rejects a -fast model id"
else
  bad "d: ca-unsafe still rejects a -fast model id" "rc=$rc out=$out"
fi

# (e) exact argv passthrough to the runtime
HOME_FOR_RUN="$FAKE_HOME"
export STUB_ARGV="$TMP/argv.bin" STUB_CGROUP="$TMP/cgroup.txt"
: >"$STUB_ARGV"
# $-bearing args are the regression case: systemd-run expands them unless escaped.
ARGS=(--print 'two words' '--flag=a b' '' '$NOT_EXPANDED' '${BRACED}' 'cost is $5 and $$x' 'tail')
run_on_tty cld-unsafe "${ARGS[@]}" >/dev/null 2>&1; rc=$?
mapfile -d '' -t GOT <"$STUB_ARGV"
if (( rc == 0 )) && [[ "${#GOT[@]}" -eq "${#ARGS[@]}" ]] && [[ "${GOT[*]}" == "${ARGS[*]}" ]]; then
  ok "e: argv passthrough is exact (${#GOT[@]} args)"
else
  bad "e: argv passthrough is exact" "rc=$rc got=$(printf '<%s>' "${GOT[@]}") want=$(printf '<%s>' "${ARGS[@]}")"
fi

# (f) one log line for an allowed invocation and one for a refused one
LOG="$FAKE_HOME/$LOG_REL"
if grep -q $'\tcld-unsafe\tallowed\t--print .*tail$' "$LOG" 2>/dev/null &&
   grep -q $'\tcld-unsafe\trefused:no-tty\t' "$LOG" 2>/dev/null &&
   grep -q $'\tcld-unsafe\trefused:agent-env-marker:AGENT_BUILD_SCOPE_ACTIVE\t' "$LOG" 2>/dev/null; then
  ok "f: audit log records allowed and refused invocations"
else
  bad "f: audit log records allowed and refused invocations" "log=$(cat "$LOG" 2>&1)"
fi
if [[ "$(head -1 "$LOG" | cut -f1)" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z$ ]]; then
  ok "f: log timestamps are ISO-8601 UTC"
else
  bad "f: log timestamps are ISO-8601 UTC" "$(head -1 "$LOG")"
fi

# (g) the invocation lands in its own transient scope under unsafe.slice
CG="$(cat "$STUB_CGROUP" 2>/dev/null)"
if [[ "$CG" == */unsafe.slice/unsafe-cld-*.scope ]]; then
  ok "g: invocation runs in its own scope under unsafe.slice ($(basename "$CG"))"
else
  bad "g: invocation runs in its own scope under unsafe.slice" "cgroup=$CG"
fi

# (h) the kernel's own view of the slice limits matches intent
SLICE_CG="/sys/fs/cgroup${CG%/*}"
read_cg() { cat "$SLICE_CG/$1" 2>/dev/null; }
swapmax=$(read_cg memory.swap.max); memmax=$(read_cg memory.max)
pidsmax=$(read_cg pids.max); cpuw=$(read_cg cpu.weight); cpumax=$(read_cg cpu.max)
printf '     kernel view of %s: memory.swap.max=%s memory.max=%s pids.max=%s cpu.weight=%s cpu.max=%s\n' \
  "$SLICE_CG" "$swapmax" "$memmax" "$pidsmax" "$cpuw" "$cpumax"
if [[ "$swapmax" == "0" ]]; then
  ok "h: memory.swap.max = 0 (swap death impossible in this slice)"
else
  bad "h: memory.swap.max = 0" "got '$swapmax'"
fi
if [[ "$memmax" == "34359738368" && "$pidsmax" == "2048" && "$cpuw" == "20" && "$cpumax" == "max 100000" ]]; then
  ok "h: memory.max=32G pids.max=2048 cpu.weight=20 and no CPU quota"
else
  bad "h: memory.max=32G pids.max=2048 cpu.weight=20 and no CPU quota" \
    "mem=$memmax pids=$pidsmax weight=$cpuw cpu.max=$cpumax"
fi

# (i) two concurrent invocations land in SEPARATE scopes
( export STUB_SLEEP=2 STUB_ARGV="$TMP/argv1.bin" STUB_CGROUP="$TMP/cg1.txt"
  run_on_tty cld-unsafe one >/dev/null 2>&1 ) &
p1=$!
( export STUB_SLEEP=2 STUB_ARGV="$TMP/argv2.bin" STUB_CGROUP="$TMP/cg2.txt"
  run_on_tty cdx-unsafe two >/dev/null 2>&1 ) &
p2=$!
wait $p1; wait $p2
cg1=$(cat "$TMP/cg1.txt" 2>/dev/null); cg2=$(cat "$TMP/cg2.txt" 2>/dev/null)
if [[ -n "$cg1" && -n "$cg2" && "$cg1" != "$cg2" &&
      "$cg1" == */unsafe.slice/*.scope && "$cg2" == */unsafe.slice/*.scope ]]; then
  ok "i: concurrent invocations get separate scopes ($(basename "$cg1") vs $(basename "$cg2"))"
else
  bad "i: concurrent invocations get separate scopes" "cg1=$cg1 cg2=$cg2"
fi

# The slice config must be the persistent one, not a shadowing user.control drop-in.
frag=$(systemctl --user show unsafe.slice -p FragmentPath --value)
drop=$(systemctl --user show unsafe.slice -p DropInPaths --value)
if [[ "$frag" == "$HOME/.config/systemd/user/unsafe.slice" && "$drop" != *user.control* ]]; then
  ok "persistence: unsafe.slice comes from ~/.config/systemd/user with no user.control shadow"
else
  bad "persistence: unsafe.slice comes from ~/.config/systemd/user with no user.control shadow" \
    "fragment=$frag dropins=$drop"
fi

printf 'agent-unsafe: pass=%d fail=%d\n' "$PASS" "$FAIL"
(( FAIL == 0 ))