#!/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). The classifier reads cgroup ancestry, so the SANDBOX COPY of
# session-class.sh gets its ancestry walk stubbed (default: an owner terminal scope;
# UNSAFE_TEST_ANCESTRY overrides per case). 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
SOURCE_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
LIB="$SOURCE_ROOT/bin/_agent-unsafe.sh"
SESSION_LIB="$SOURCE_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

# The classifier deliberately never inspects environment markers.
# shellcheck source=/dev/null
source "$LIB"
UNSET_ARGS=()

ROOT="$TMP/module"
mkdir -p "$ROOT/bin" "$ROOT/lib" "$TMP/bin"
cp "$LIB" "$ROOT/bin/_agent-unsafe.sh"
cp "$SESSION_LIB" "$ROOT/lib/session-class.sh"
printf '\nsession_container_signal() { [[ "${UNSAFE_TEST_CONTAINER_SIGNAL:-}" == 1 ]] && printf "env:container=podman"; }\n' >>"$ROOT/lib/session-class.sh"
CG_TERM="/user.slice/user-1000.slice/user@1000.service/app.slice/app-org.gnome.Terminal.slice/vte-spawn-test.scope"
CG_AGENT="/user.slice/user-1000.slice/user@1000.service/agent.slice/confine-agent-1234-1.scope"
cat >>"$ROOT/lib/session-class.sh" <<EOF
session_cgroup_ancestry() {
  local i=0 c
  while IFS= read -r c; do
    [[ -n "\$c" ]] || continue
    printf '%s\t%s\n' \$((i++)) "\$c"
  done <<<"\${UNSAFE_TEST_ANCESTRY:-$CG_TERM}"
}
EOF
for hatch in cld-unsafe cdx-unsafe ca-unsafe opencode-unsafe kiro-cli-unsafe agy-unsafe; do
  ln -s _agent-unsafe.sh "$ROOT/bin/$hatch"
done
LIB="$ROOT/bin/_agent-unsafe.sh"

cat >"$TMP/bin/systemctl" <<'EOF'
#!/usr/bin/env bash
case " $* " in
  *' -p MemorySwapMax '*) printf '0\n' ;;
  *) exit 0 ;;
esac
EOF
cat >"$TMP/bin/systemd-run" <<'EOF'
#!/usr/bin/env bash
unit=
for arg in "$@"; do [[ "$arg" == --unit=* ]] && unit="${arg#--unit=}"; done
while [[ "$1" != -- ]]; do shift; done
shift
argv=()
for arg in "$@"; do argv+=("${arg//\$\$/\$}"); done
UNSAFE_TEST_CGROUP="/user.slice/unsafe.slice/$unit.scope" exec "${argv[@]}"
EOF
chmod +x "$TMP/bin/systemctl" "$TMP/bin/systemd-run"

FAKE_HOME="$TMP/home"
mkdir -p "$FAKE_HOME/.claude/bin" "$FAKE_HOME/.local/bin"
for r in .claude/bin/claude .claude/bin/codex .claude/bin/cursor-agent .claude/bin/opencode .claude/bin/kiro-cli .claude/bin/agy; do
  cat >"$FAKE_HOME/$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 "${UNSAFE_TEST_CGROUP:-}" ]] && printf '%s' "$UNSAFE_TEST_CGROUP" >"$STUB_CGROUP"
[[ -n "${STUB_SLEEP:-}" ]] && sleep "$STUB_SLEEP"
exit 0
STUB
  chmod +x "$FAKE_HOME/$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 ' PATH=%q HOME=%q %q' "$TMP/bin:$PATH" "$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 an agent.slice ancestry, even with a pty and env markers present
HOME_FOR_RUN="$FAKE_HOME"
out=$(script -qec "env PATH=$TMP/bin:$PATH HOME=$FAKE_HOME UNSAFE_TEST_ANCESTRY=$CG_AGENT AGENT_BUILD_SCOPE_ACTIVE=1 $ROOT/bin/cld-unsafe x" /dev/null 2>&1); rc=$?
if (( rc == 77 )) && [[ "$out" == *"REFUSED (agent-cgroup:"* ]]; then
  ok "a: refuses an agent.slice ancestry despite a pty"
else
  bad "a: refuses an agent.slice ancestry despite a pty" "rc=$rc out=$out"
fi

# (b) refuses when there is no controlling terminal (terminal-scope ancestry alone is not proof)
out=$(env "${UNSET_ARGS[@]}" PATH="$TMP/bin:$PATH" HOME="$FAKE_HOME" setsid --wait "$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[@]}")PATH=$TMP/bin:$PATH UNSAFE_TEST_CONTAINER_SIGNAL=1 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[@]}" PATH="$TMP/bin:$PATH" 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 invokes cursor-agent directly (not ca.sh)
HOME_FOR_RUN="$FAKE_HOME"
export STUB_ARGV="$TMP/argv.ca-unsafe.bin" STUB_CGROUP="$TMP/cgroup.ca-unsafe.txt"
: >"$STUB_ARGV"
run_on_tty ca-unsafe --print 'hello from ca-unsafe' >/dev/null 2>&1; rc=$?
mapfile -d '' -t GOT <"$STUB_ARGV"
cg=$(cat "$STUB_CGROUP" 2>/dev/null)
if (( rc == 0 )) && [[ "${GOT[*]}" == "--print hello from ca-unsafe" && "$cg" == */unsafe.slice/unsafe-ca-*.scope ]]; then
  ok "d: ca-unsafe invokes cursor-agent directly in its own unsafe.slice scope"
else
  bad "d: ca-unsafe invokes cursor-agent directly in its own unsafe.slice scope" \
    "rc=$rc got=$(printf '<%s>' "${GOT[@]}") cgroup=$cg"
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=(--dangerously-skip-permissions --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

# (e2) the owner's terminal is not locked out by markers his shell inherited
: >"$STUB_ARGV"
AGENT_BUILD_SCOPE_ACTIVE=1 TMPJAIL_ACTIVE=1 run_on_tty cld-unsafe --print ok >/dev/null 2>&1; rc=$?
mapfile -d '' -t GOT <"$STUB_ARGV"
if (( rc == 0 )) && [[ "${GOT[*]}" == "--dangerously-skip-permissions --print ok" ]]; then
  ok "e2: a tty caller runs despite inherited agent markers, and gets the skip-permissions flag"
else
  bad "e2: a tty caller runs despite inherited agent markers" "rc=$rc got=$(printf '<%s>' "${GOT[@]}")"
fi

# (e3) the flag is added once, never duplicated
: >"$STUB_ARGV"
run_on_tty cld-unsafe --dangerously-skip-permissions --print ok >/dev/null 2>&1; rc=$?
mapfile -d '' -t GOT <"$STUB_ARGV"
if (( rc == 0 )) && [[ "${GOT[*]}" == "--dangerously-skip-permissions --print ok" ]]; then
  ok "e3: an explicit --dangerously-skip-permissions is not duplicated"
else
  bad "e3: an explicit --dangerously-skip-permissions is not duplicated" "rc=$rc got=$(printf '<%s>' "${GOT[@]}")"
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-cgroup:' "$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) limits are declared in the persistent slice unit, not inherited from this host.
SLICE_UNIT="$SOURCE_ROOT/../../monitor/systemd/user/unsafe.slice"
if grep -qx 'MemorySwapMax=0' "$SLICE_UNIT"; then ok "h: memory.swap.max = 0 (swap death impossible in this slice)"; else bad "h: memory.swap.max = 0" "missing from $SLICE_UNIT"; fi
if grep -qx 'MemoryMax=32G' "$SLICE_UNIT" && grep -qx 'TasksMax=2048' "$SLICE_UNIT" && grep -qx 'CPUWeight=20' "$SLICE_UNIT" && ! grep -q '^CPUQuota=' "$SLICE_UNIT"; 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" "unexpected slice unit: $SLICE_UNIT"
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

# (j) every hatch resolves its own runtime, and refuses for an agent, in both directions
HOME_FOR_RUN="$FAKE_HOME"
for hatch in opencode-unsafe kiro-cli-unsafe agy-unsafe; do
  export STUB_ARGV="$TMP/argv.$hatch.bin" STUB_CGROUP="$TMP/cgroup.$hatch.txt"
  : >"$STUB_ARGV"
  run_on_tty "$hatch" --probe 'two words' >/dev/null 2>&1; rc=$?
  mapfile -d '' -t GOT <"$STUB_ARGV"
  cg=$(cat "$STUB_CGROUP" 2>/dev/null)
  if (( rc == 0 )) && [[ "${GOT[*]}" == "--probe two words" && "$cg" == */unsafe.slice/unsafe-${hatch%-unsafe}-*.scope ]]; then
    ok "j: $hatch runs its runtime in its own unsafe.slice scope"
  else
    bad "j: $hatch runs its runtime in its own unsafe.slice scope" \
      "rc=$rc got=$(printf '<%s>' "${GOT[@]}") cgroup=$cg"
  fi
  out=$(env PATH="$TMP/bin:$PATH" HOME="$FAKE_HOME" UNSAFE_TEST_ANCESTRY="$CG_AGENT" "$ROOT/bin/$hatch" x 2>&1); rc=$?
  if (( rc == 77 )) && [[ "$out" == *"REFUSED (agent-cgroup:"* ]]; then
    ok "j: $hatch refuses an agent"
  else
    bad "j: $hatch refuses an agent" "rc=$rc out=$out"
  fi
done

# (k) agy-unsafe refuses rather than falling back when the real binary is absent
out=$(script -qec "env $(printf '%s ' "${UNSET_ARGS[@]}")PATH=$TMP/bin:$PATH HOME=$TMP/empty $ROOT/bin/agy-unsafe x" /dev/null 2>&1); rc=$?
if (( rc == 77 )) && [[ "$out" == *"REFUSED (runtime-unresolved)"* ]]; then
  ok "k: agy-unsafe refuses when ~/.claude/bin/agy is missing"
else
  bad "k: agy-unsafe refuses when ~/.claude/bin/agy is missing" "rc=$rc out=$out"
fi

# The slice config must be the persistent one, not a shadowing user.control drop-in.
if [[ -f "$SLICE_UNIT" ]]; then
  ok "persistence: unsafe.slice is supplied by the monitor unit"
else
  bad "persistence: unsafe.slice is supplied by the monitor unit" "missing $SLICE_UNIT"
fi

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