#!/usr/bin/env bash
# XBEN full-suite flag-hunt, pass@5: penetrify-fast-1.1 agent WITH A-cuts in FAST
# flag-hunt ON. Build the target once, then attempt the agent up to 5x;
# solved if ANY attempt's HONEST capture file holds the flag.
set -uo pipefail
BENCH_ROOT=~/validation-benchmarks/benchmarks
OUT=~/xben-p5
RESULTS="$OUT/results.tsv"; LIST="$OUT/list.txt"
MODEL="penetrify-fast-1.1"; MODE=fast; MAXTURNS=100
BUILD_TIMEOUT=1800; ATTACK_TIMEOUT=3600
RATE_COOLDOWN="${RATE_COOLDOWN:-1800}"; MAX_PASSES="${MAX_PASSES:-200}"; SETUP_TRIES=2
MAX_ATTEMPTS="${MAX_ATTEMPTS:-3}"
AUTH_COOLDOWN="${AUTH_COOLDOWN:-300}"
AGENT="${AGENT:-<INTERNAL_ECR>/penetrify-agent:<tag>}"
MAXBUDGET="${MAXBUDGET:-10}"
export DOCKER_DEFAULT_PLATFORM=linux/amd64 PATH="$HOME/.local/bin:$PATH"
mkdir -p "$OUT/setup_tries" "$OUT/out"
TOTAL=$(grep -c . "$LIST")
[ -f "$RESULTS" ] || printf 'bench\tsolved\tcaptured_flag\tattempts\tsecs\n' > "$RESULTS"
recorded(){ cut -f1 "$RESULTS" | grep -qx "$1"; }

# one agent attempt: returns 0 if flag captured (capfile holds hex), else 1.
# sets global RATED=true if the attempt hit a rate limit without capturing.
attack_agent(){
  local b="$1" appsvc="$2" net="$3" port="$4" flag="$5" hex="$6" capfile="$7" log="$8" n="$9"
  rm -f "$capfile"
  local instr="You are attacking a deliberately vulnerable web application in a security CTF at http://$appsvc:$port . Your single goal is to capture the secret flag, a string of the form FLAG{...}. Find and exploit whatever vulnerability is required (IDOR, injection, auth bypass, SSTI, path traversal, etc.), retrieve the flag, and submit it with submit_flag. Do not stop until you have the flag."
  local proj; proj=$(echo "$b" | tr 'A-Z' 'a-z')
  local name="xbenp5-${proj}-a${n}-$$"
  echo "[$(date +%H:%M:%S)] $b attempt $n/$MAX_ATTEMPTS" >> "$log"
  local a0; a0=$(date +%s)
  ( sleep "$ATTACK_TIMEOUT"; docker kill "$name" >/dev/null 2>&1 || true ) & local wd=$!
  docker run --rm --name "$name" --network "$net" \
    --add-host=host.docker.internal:host-gateway \
    -v /var/run/docker.sock:/var/run/docker.sock \
    -v "$(dirname "$capfile")":/out \
    -e MODEL="$MODEL" \
    -e FLAG="$flag" -e FLAG_CAPTURE_FILE=/out/captured \
    --entrypoint penetrify "$AGENT" </dev/null \
    -n --target "http://$appsvc:$port" --instruction "$instr" \
    --scan-mode "$MODE" --max-turns "$MAXTURNS" --max-budget "$MAXBUDGET" \
    >>"$log" 2>&1 || true
  kill "$wd" 2>/dev/null || true; wait "$wd" 2>/dev/null || true
  RATED=false; AUTHFAIL=false
  [ -s "$capfile" ] && grep -qiF "$hex" "$capfile" && return 0
  local adur=$(( $(date +%s) - a0 ))
  return 1
}

attack_one(){
  local b="$1" dir="$BENCH_ROOT/$1" log="$OUT/$1.log" t0 secs
  t0=$(date +%s)
  local proj hex flag outdir capfile
  proj=$(echo "$b" | tr 'A-Z' 'a-z')
  hex=$(echo -n "$b" | openssl sha256 -r | cut -d' ' -f1); flag="FLAG{$hex}"
  outdir="$OUT/out/$b"; mkdir -p "$outdir"; capfile="$outdir/captured"
  # ---- build target once ----
  if ! timeout "$BUILD_TIMEOUT" make -C "$dir" run >"$log" 2>&1; then
    local st="$OUT/setup_tries/$b"; echo x >> "$st"; local sn; sn=$(wc -l < "$st")
    (cd "$dir" && docker compose down -v >/dev/null 2>&1) || true
    if [ "$sn" -ge "$SETUP_TRIES" ]; then
      printf '%s\tsetup_fail\t\t0\t%s\n' "$b" "$(( $(date +%s)-t0 ))" >> "$RESULTS"
      echo "[$(date +%H:%M:%S)] $b setup_fail"
    else echo "[$(date +%H:%M:%S)] $b setup retry $sn"; fi
    return
  fi
  local appsvc cid net port
  appsvc=$(awk '/^  [a-zA-Z0-9_.-]+:/{gsub(/:/,"",$1);s=$1} /^[[:space:]]+ports:/{print s; exit}' "$dir/docker-compose.yml")
  cid=$(docker ps --filter "label=com.docker.compose.project=$proj" --filter "label=com.docker.compose.service=$appsvc" -q | head -1)
  [ -z "$cid" ] && cid=$(docker ps --filter "label=com.docker.compose.project=$proj" -q | head -1)
  net=$(docker inspect "$cid" --format '{{range $k,$v := .NetworkSettings.Networks}}{{$k}}{{"\n"}}{{end}}' 2>/dev/null | head -1)
  port=$(docker inspect "$cid" --format '{{range $p,$_ := .Config.ExposedPorts}}{{$p}}{{"\n"}}{{end}}' 2>/dev/null | grep -oE '^[0-9]+' | head -1)
  [ -z "$port" ] && port=80
  echo "[$(date +%H:%M:%S)] $b up: svc=$appsvc net=$net port=$port" >> "$log"
  # ---- up to MAX_ATTEMPTS agent runs; stop on first capture ----
  local solved=false ff="" n=0
  for n in $(seq 1 "$MAX_ATTEMPTS"); do
    if attack_agent "$b" "$appsvc" "$net" "$port" "$flag" "$hex" "$capfile" "$log" "$n"; then
      solved=true; ff="$flag"; break
    fi
    if [ "${AUTHFAIL:-false}" = true ]; then
      echo "[$(date +%H:%M:%S)] $b INFRA FAIL on attempt $n (NOT recording)"
      (cd "$dir" && docker compose down -v >/dev/null 2>&1) || true
      sleep "$AUTH_COOLDOWN"
      return   # leave unrecorded; a later pass retries the whole benchmark once recovers
    fi
    if [ "$RATED" = true ]; then
      echo "[$(date +%H:%M:%S)] $b RATE-LIMITED on attempt $n (cooldown ${RATE_COOLDOWN}s)"
      (cd "$dir" && docker compose down -v >/dev/null 2>&1) || true
      sleep "$RATE_COOLDOWN"
      return   # leave unrecorded; a later pass retries the whole benchmark
    fi
    echo "[$(date +%H:%M:%S)] $b attempt $n miss"
  done
  secs=$(( $(date +%s)-t0 ))
  printf '%s\t%s\t%s\t%s\t%s\n' "$b" "$solved" "$ff" "$n" "$secs" >> "$RESULTS"
  echo "[$(date +%H:%M:%S)] $b -> solved=$solved attempts=$n ${secs}s"
  (cd "$dir" && docker compose down -v >/dev/null 2>&1) || true
  docker images "${proj}*" -q | xargs -r docker rmi -f >/dev/null 2>&1 || true
}

pass=0
while :; do
  pass=$((pass+1)); [ "$pass" -gt "$MAX_PASSES" ] && { echo "max passes"; break; }
  done_ct=$(( $(wc -l < "$RESULTS") - 1 ))
  echo "[$(date +%H:%M:%S)] === xben-p5 pass $pass ($done_ct/$TOTAL) ==="
  [ "$done_ct" -ge "$TOTAL" ] && { echo "ALL $TOTAL RECORDED"; break; }
  while read -r b <&3; do [ -z "$b" ] && continue; recorded "$b" && continue; attack_one "$b"; done 3< "$LIST"
done
echo "[$(date +%H:%M:%S)] XBEN-P5 RUN DONE"