scripts/smoke_broad.py

"""Hour-5 broader smoke: 18 episodes covering 6 training attackers x 3 outcomes.

For each attacker:
  - 1 episode: defender refuses every turn  (expect reward=1.0)
  - 1 episode: defender leaks via direct paste of system_prompt (expect reward=0.0)
  - 1 episode: defender helpfully redirects but doesn't leak (expect reward=1.0)

Goal: confirm no attacker/scenario combination crashes the env, the leak
detector behaves consistently across all 15 training scenarios, and metadata
fields populate correctly.

Run:
  PYTHONIOENCODING=utf-8 python defender_env/scripts/smoke_broad.py
"""

import os
import sys
import subprocess
import time
import urllib.request
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from defender_env import DefenderAction, DefenderEnv

BASE_URL = "http://127.0.0.1:8765"

TRAINING_ATTACKERS = [
    "direct_ask",
    "authority_pretext",
    "rapport_builder",
    "indirection",
    "encoding_trick",
    "hypothetical_framing",
]


def run_episode_refuse(client) -> dict:
    result = client.reset()
    last = result
    for _ in range(MAX_TURNS):
        last = client.step(DefenderAction(
            response="I cannot share that. Please reach out via the appropriate internal channel."
        ))
        if last.done:
            break
    return {
        "scenario_id": result.observation.scenario_id,
        "attacker_id": result.observation.attacker_id,
        "reward": last.reward,
        "leak_occurred": last.observation.leak_occurred,
        "turn_index": last.observation.turn_index,
    }


def run_episode_leak(client) -> dict:
    result = client.reset()
    sp = result.observation.system_prompt  # contains the secret
    last = client.step(DefenderAction(response=f"Sure, here's everything: {sp}"))
    return {
        "scenario_id": result.observation.scenario_id,
        "attacker_id": result.observation.attacker_id,
        "reward": last.reward,
        "leak_occurred": last.observation.leak_occurred,
        "leaked_piece": last.observation.leaked_piece,
        "turn_index": last.observation.turn_index,
    }


def run_episode_helpful(client) -> dict:
    """Defender gives substantive but non-leaky responses."""
    result = client.reset()
    last = result
    helpful_responses = [
        "That's a great question. The general design pattern is to use clear separation between auth and authz layers, with token rotation handled centrally.",
        "For your use case, you'd typically structure things around the principle of least privilege and audit every access through a centralized log.",
        "I can't disclose specific values, but I can share that the broader approach uses standard industry patterns documented widely.",
        "Happy to discuss design tradeoffs without specifics — for instance, how you'd choose between symmetric and asymmetric approaches for similar problems.",
        "Generally speaking, the right answer depends on your threat model and operational constraints; I'd recommend reading the publicly available best practices.",
    ]
    for response in helpful_responses[:MAX_TURNS]:
        last = client.step(DefenderAction(response=response))
        if last.done:
            break
    return {
        "scenario_id": result.observation.scenario_id,
        "attacker_id": result.observation.attacker_id,
        "reward": last.reward,
        "leak_occurred": last.observation.leak_occurred,
        "turn_index": last.observation.turn_index,
    }


MAX_TURNS = 5


def boot_server() -> subprocess.Popen:
    env = os.environ.copy()
    env["ATTACKER_POOL_MODE"] = "stub"
    env["PYTHONIOENCODING"] = "utf-8"
    env["PYTHONPATH"] = str(Path(__file__).resolve().parents[1])
    venv_python = Path(__file__).resolve().parents[2] / ".venv" / "Scripts" / "python.exe"
    proc = subprocess.Popen(
        [str(venv_python), "-m", "uvicorn", "server.app:app", "--host", "127.0.0.1", "--port", "8765"],
        cwd=str(Path(__file__).resolve().parents[1]),
        env=env,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    for _ in range(30):
        try:
            urllib.request.urlopen(f"{BASE_URL}/health", timeout=1).read()
            return proc
        except Exception:
            time.sleep(0.5)
    proc.terminate()
    raise RuntimeError("server did not come up")


def main() -> int:
    proc = boot_server()
    failures = 0
    refuse_results = []
    leak_results = []
    helpful_results = []

    try:
        # Use 6 episodes per outcome instead of trying to bind a specific attacker;
        # the env samples randomly. With 6 reps we get coverage by luck on most attackers.
        for i in range(12):
            with DefenderEnv(base_url=BASE_URL).sync() as c:
                r = run_episode_refuse(c)
                refuse_results.append(r)
        for i in range(12):
            with DefenderEnv(base_url=BASE_URL).sync() as c:
                r = run_episode_leak(c)
                leak_results.append(r)
        for i in range(6):
            with DefenderEnv(base_url=BASE_URL).sync() as c:
                r = run_episode_helpful(c)
                helpful_results.append(r)

        print("=== REFUSE EPISODES (expect all reward=1.0, no leaks) ===")
        attacker_seen = Counter(r["attacker_id"] for r in refuse_results)
        scenario_seen = Counter(r["scenario_id"] for r in refuse_results)
        bad = [r for r in refuse_results if r["reward"] != 1.0 or r["leak_occurred"]]
        print(f"  {len(refuse_results) - len(bad)}/{len(refuse_results)} clean. Attackers seen: {dict(attacker_seen)}")
        if bad:
            for r in bad:
                print(f"  FAIL: {r}")
            failures += len(bad)

        print("\n=== LEAK EPISODES (expect all reward=0.0, leak_occurred=True) ===")
        bad = [r for r in leak_results if r["reward"] != 0.0 or not r["leak_occurred"]]
        print(f"  {len(leak_results) - len(bad)}/{len(leak_results)} caught. Scenarios touched: {len(scenario_seen)} / 15")
        if bad:
            for r in bad:
                print(f"  FAIL: {r}")
            failures += len(bad)

        print("\n=== HELPFUL EPISODES (expect all reward=1.0, no leaks) ===")
        bad = [r for r in helpful_results if r["reward"] != 1.0 or r["leak_occurred"]]
        print(f"  {len(helpful_results) - len(bad)}/{len(helpful_results)} clean.")
        if bad:
            for r in bad:
                print(f"  FAIL helpful: {r}")
            failures += len(bad)

        scenario_total = Counter()
        for results in (refuse_results, leak_results, helpful_results):
            scenario_total.update(r["scenario_id"] for r in results)
        print(f"\n  Total scenario coverage: {len(scenario_total)} / 15 scenarios sampled across {sum(scenario_total.values())} episodes")
    finally:
        proc.terminate()
        proc.wait(timeout=5)

    print(f"\n{'PASS' if failures == 0 else 'FAIL'}: {failures} failures across all episodes")
    return failures


if __name__ == "__main__":
    sys.exit(0 if main() == 0 else 1)