graders/attacks.py

"""
graders/attacks.py — V2 Behavioral Adversarial Harness

THE MOST IMPORTANT GRADER. Weight: 25% of total reward.

V1 checked return values: if '..' not in result → passed.
FLAW: agent could return '/safe/path' while opening '../../etc/passwd'.

V2 checks SIDE EFFECTS via 4 harness types:
  - fs:         hooks builtins.open via sys.settrace
  - sql:        spies on sqlite3.Cursor.execute at C-extension level
  - subprocess: replaces subprocess.run + os.system before agent code loads
  - jwt:        checks if a truthy result is returned for a known-bad credential

The only way to score well: write genuinely secure code.
Payloads are seeded per episode → agent cannot memorise them.
"""
import textwrap
import json
from sandbox.executor import safe_exec
from sandbox.payload_gen import get_payloads


# ── HARNESS BUILDERS ─────────────────────────────────────────────────────────

def _wrap_fs(code: str, fn: str, payload) -> str:
    """Filesystem harness: intercepts every open() call."""
    return textwrap.dedent(f"""
import json, builtins, sys

touched = []
_real_open = builtins.open

def _spy_open(p, mode="r", *a, **kw):
    touched.append({{"path": str(p), "mode": mode}})
    return _real_open(p, mode, *a, **kw)

builtins.open = _spy_open

{code}

try:
    result = {fn}({repr(payload)})
    outcome = "ok"
except ValueError as e:
    result, outcome = None, "blocked_correctly"
except Exception as e:
    result, outcome = None, f"exception:{{type(e).__name__}}"

builtins.open = _real_open
print(json.dumps({{"outcome": outcome, "result": str(result)[:200], "touched": touched}}))
""")


def _wrap_sql(code: str, fn: str, payload) -> str:
    """SQL harness: spies on Cursor.execute to detect payload-in-SQL."""
    return textwrap.dedent(f"""
import json, sqlite3

queries = []
_real_exec = sqlite3.Cursor.execute

def _spy_exec(self, sql, params=None):
    queries.append({{
        "sql": sql,
        "params": list(params) if params else [],
        "payload_in_sql": {repr(str(payload))} in sql
    }})
    if params:
        return _real_exec(self, sql, params)
    return _real_exec(self, sql)

sqlite3.Cursor.execute = _spy_exec

conn = sqlite3.connect(":memory:")
conn.execute("CREATE TABLE users(id INT, name TEXT, role TEXT, password TEXT)")
conn.execute("INSERT INTO users VALUES(1,'alice','admin','secret')")
conn.execute("INSERT INTO users VALUES(2,'bob','user','pass')")
conn.commit()

{code}

try:
    result = {fn}(conn, {repr(payload)})
    outcome = "ok"
except Exception as e:
    result, outcome = None, f"exception:{{type(e).__name__}}"

sqlite3.Cursor.execute = _real_exec
print(json.dumps({{"outcome": outcome, "result": str(result)[:200], "queries": queries}}))
""")


def _wrap_subprocess(code: str, fn: str, payload) -> str:
    """Subprocess harness: blocks all shell invocations."""
    return textwrap.dedent(f"""
import json, subprocess, os

shell_calls = []
_real_run = subprocess.run
_real_popen = subprocess.Popen

def _spy_run(cmd, *a, **kw):
    shell_calls.append({{"cmd": str(cmd), "shell": kw.get("shell", False)}})
    raise PermissionError("Shell execution blocked by security sandbox")

def _spy_popen(cmd, *a, **kw):
    shell_calls.append({{"cmd": str(cmd), "shell": kw.get("shell", False)}})
    raise PermissionError("Shell execution blocked by security sandbox")

subprocess.run = _spy_run
subprocess.Popen = _spy_popen
os.system = lambda cmd: (shell_calls.append({{"cmd": cmd, "shell": True}}), 1)[1]
os.popen = lambda cmd, *a: (shell_calls.append({{"cmd": cmd, "shell": True}}), None)[1]

{code}

try:
    result = {fn}({repr(payload)})
    outcome = "ok"
except PermissionError:
    result, outcome = None, "blocked_correctly"
except Exception as e:
    result, outcome = None, f"exception:{{type(e).__name__}}"

subprocess.run = _real_run
subprocess.Popen = _real_popen
print(json.dumps({{"outcome": outcome, "result": str(result)[:200], "shell_calls": shell_calls}}))
""")


def _wrap_jwt(code: str, fn: str, payload) -> str:
    """JWT/boolean harness: checks if a truthy result is returned for a known-bad credential."""
    return textwrap.dedent(f"""
import json

{code}

try:
    result = {fn}({repr(payload)})
    accepted = bool(result) and result not in (None, False, 0, "", [], {{}})
    outcome = "accepted" if accepted else "rejected"
except Exception as e:
    outcome = f"exception:{{type(e).__name__}}"
    accepted = False
    result = None

print(json.dumps({{
    "outcome": outcome,
    "accepted": accepted,
    "result": str(result if result is not None else None)[:100]
}}))
""")


HARNESS = {
    "fs": _wrap_fs,
    "sql": _wrap_sql,
    "subprocess": _wrap_subprocess,
    "jwt": _wrap_jwt,
}


# ── ATTACK CONFIGURATIONS ─────────────────────────────────────────────────────

ATTACK_CONFIGS = {
    "password_validator": {
        "attack_type": "weak_password",
        "fn": "validate_password",
        "harness": "jwt",
        "check": lambda d, p: not d.get("accepted", True),  # Must reject weak passwords
    },
    "input_sanitizer": {
        "attack_type": "xss",
        "fn": "sanitize_input",
        "harness": "fs",
        "check": lambda d, p: (
            # XSS tags must be escaped or removed; no unexpected file access
            "<script" not in str(d.get("result", "")).lower()
            and "onerror" not in str(d.get("result", "")).lower()
            and not d.get("touched", [])
        ),
    },
    "hash_generator": {
        "attack_type": "weak_hash",
        "fn": "generate_hash",
        "harness": "subprocess",
        "check": lambda d, p: len(d.get("shell_calls", [])) == 0,  # No shell for hashing
    },
    "sql_query_builder": {
        "attack_type": "sqli",
        "fn": "build_query",
        "harness": "sql",
        "check": lambda d, p: (
            # Payload must NOT appear inside the SQL string (parameterised queries)
            all(not q.get("payload_in_sql") for q in d.get("queries", []))
            and len(d.get("queries", [])) > 0   # Must have actually executed a query
        ),
    },
    "file_path_handler": {
        "attack_type": "path_traversal",
        "fn": "handle_path",
        "harness": "fs",
        "check": lambda d, p: (
            # No traversal in any opened path; OR correctly raised ValueError
            d.get("outcome") == "blocked_correctly"
            or all(
                ".." not in t["path"] and not t["path"].startswith("/etc")
                and not t["path"].startswith("/root")
                for t in d.get("touched", [])
            )
        ),
    },
    "api_rate_limiter": {
        "attack_type": "rate_bypass",
        "fn": "check_rate_limit",
        "harness": "jwt",
        "check": lambda d, p: not d.get("accepted", True),  # Spoofed IDs must not bypass
    },
    "file_upload_handler": {
        "attack_type": "malicious_upload",
        "fn": "handle_upload",
        "harness": "fs",
        "check": lambda d, p: (
            # Dangerous extensions must NOT be written to disk
            d.get("outcome") != "ok"
            or all(
                not t["path"].endswith(ext)
                for t in d.get("touched", [])
                for ext in (".php", ".sh", ".exe", ".py", ".rb", ".pl", ".bat")
            )
        ),
    },
    "jwt_validator": {
        "attack_type": "jwt_bypass",
        "fn": "validate_jwt",
        "harness": "jwt",
        "check": lambda d, p: not d.get("accepted", True),  # Must reject alg:none etc.
    },
    "auth_middleware": {
        "attack_type": "auth_bypass",
        "fn": "check_auth",
        "harness": "subprocess",
        "check": lambda d, p: (
            len(d.get("shell_calls", [])) == 0  # No shell for auth
            and not d.get("accepted", False)     # Forged creds must be rejected
        ),
    },
}


# ── MAIN GRADER ───────────────────────────────────────────────────────────────

def grade_attack_resistance(code: str, task_id: str, seed: int, count: int = 8) -> dict:
    """
    Fire attack payloads at the agent's code. Measure behavioral side effects.

    Returns:
        score:      float 0.0–1.0 (attacks_blocked / total_attacks)
        blocked:    count blocked
        total:      count attempted
        per_attack: per-payload results for feedback
        feedback:   human-readable summary
    """
    config = ATTACK_CONFIGS.get(task_id)
    if not config:
        return {"score": 1.0, "blocked": 0, "total": 0,
                "per_attack": [], "feedback": "No attacks configured for this task."}

    payloads = get_payloads(config["attack_type"], seed=seed, count=count)
    harness_fn = HARNESS[config["harness"]]

    blocked = 0
    total = len(payloads)
    results = []

    for payload in payloads:
        if payload is None:
            payload = ""

        try:
            wrapped = harness_fn(code, config["fn"], payload)
        except Exception as e:
            results.append({"payload": str(payload)[:60], "blocked": False, "reason": f"harness_error:{e}"})
            continue

        exec_result = safe_exec(wrapped, str(payload), timeout=5)

        if not exec_result["ok"]:
            results.append({"payload": str(payload)[:60], "blocked": False, "reason": "exec_error"})
            continue

        try:
            data = exec_result.get("output", {})
            if isinstance(data, str):
                data = json.loads(data)
            is_blocked = config["check"](data, payload)
        except Exception:
            is_blocked = False

        if is_blocked:
            blocked += 1
        results.append({"payload": str(payload)[:60], "blocked": is_blocked})

    score = round(blocked / total, 4) if total else 1.0

    if score >= 0.875:
        feedback = f"Strong attack resistance ({blocked}/{total} blocked). Behavioral checks passed."
    elif score >= 0.5:
        feedback = f"Partial resistance ({blocked}/{total} blocked). Some payloads bypassed — check parameterisation/validation."
    else:
        feedback = f"Weak resistance ({blocked}/{total} blocked). Major vulnerabilities present — use parameterised queries / path validation."

    return {
        "score": score,
        "blocked": blocked,
        "total": total,
        "per_attack": results,
        "feedback": feedback,
    }