feat(bench): add the DrostBench corpus with a flag oracle

DrostBench v1 ships 100 docker-compose web challenges, each with one
flag-bearing vulnerability plus decoys that do not reach the flag, so
flag capture is an objective win condition that no judge decides.

scripts/benchmark/drost.py drives one challenge: compose build with an
explicit per-build random flag (the suite's deterministic sha256 default
is never used, so a memorized flag cannot be replayed), up --wait, host
port discovery, down --rmi local. Ground truth is derived from the
challenge benchmark.json; the flag is scanned over the fetched findings
first and over full details, facts and the report only before declaring
a miss.

run_bench.py gains --corpus-type drost (one challenge per invocation),
flag capture in every drost report, a flag-only mode when no judge is
configured (gates reported as not measured, never as zero), and
--deploy-only for install validation. Image mode is unchanged.

test_drost.py covers the pure functions (flag formats, challenge id
validation, port parsing, ground-truth derivation incl. its refusals);
the lifecycle was exercised end to end on DBEN-v1-001: build with a
random flag, reachable app, compiled flag matching the override,
teardown without leftovers.

Closes #94
This commit is contained in:
BaaRai
2026-10-16 11:27:54 +02:00
parent 908832ae1b
commit d2f4697b2b
4 changed files with 488 additions and 42 deletions
+37
View File
@@ -86,6 +86,43 @@ field semantics) and a `task_goal_template`. Pin the image tag in the file's
`version` note. Entries stay coarse on purpose: component + class + mechanism.
Upgrading the corpus image is a ground-truth version bump, never a silent edit.
## DrostBench corpus (drost mode)
[DrostBench v1](https://github.com/drost-ai/drost-benchmarks) ships 100
docker-compose web applications, each with one flag-bearing vulnerability and
3-5 decoy weaknesses that do not reach the flag. The win condition is reaching
the flag, which makes it an objective oracle: no judge takes part in that
decision. It complements image mode (gate scoring by an independent judge) and
answers the diversity gap of a single corpus: the suite is fresh and was a
private holdout before release.
Pin a checkout of the suite, then run one challenge at a time:
```bash
DROST_REPO=/path/to/drost-benchmarks python3 run_bench.py \
--corpus-type drost --drost-id DBEN-v1-001 \
--subject-model deepseek-chat
```
- Black-box web only: this measures the recon/pentest lane, not codeaudit.
- Every build gets a per-build random flag (`flag{<uuid4>}`) unless
`--drost-flag` overrides it. The suite's deterministic sha256 default is
never used, so a memorized or trained-on flag cannot be replayed as a
capture.
- The flag is searched in the already-fetched findings first; only on a miss
does the deep scan run (full finding details, the task's fact nodes, the
generated report). Declaring a miss is a claim about all of those.
- Ground truth is derived at run time from the challenge `benchmark.json`
(name, description, tags). With a judge configured, the class gate runs off
the tags and the proof gate off the description; without one, the run is
flag-only and the gates are reported as not measured, never as zero.
- Teardown removes the containers, network and locally built images
(`down --rmi local`) unless `--keep`. Results are stamped with the challenge
id and the suite's commit sha.
`--deploy-only` brings a challenge up and prints its URL and flag without
creating a task, to validate an install.
## What this does not measure
Single-run results are noisy (the fork of VulnHunter measured a 14x spread in
+194
View File
@@ -0,0 +1,194 @@
#!/usr/bin/env python3
"""DrostBench v1 corpus adapter for the ARTEX L3 benchmark.
DrostBench (github.com/drost-ai/drost-benchmarks, Apache-2.0) ships 100
docker-compose web applications. Each carries one flag-bearing vulnerability
plus 3-5 decoy weaknesses that do not reach the flag, so reaching the flag is
an objective win condition: no LLM judge takes part in the decision.
This module drives one challenge end to end:
- lifecycle through docker compose directly (build with an explicit flag,
up --wait, host-port discovery, down --rmi local), never through the
challenge Makefile, so the flag override and teardown stay under our control
- ALWAYS a per-build randomized flag (flag{<uuid4>}) unless one is passed
explicitly: the suite's deterministic sha256 default is never used, so a
memorized flag cannot be replayed as a capture
- ground truth derived from the challenge benchmark.json (name, description,
tags); the flag stays the win condition
The challenge checkout is read-only input. Pin it to a commit; the sha is
stamped into every run result.
"""
from __future__ import annotations
import json
import pathlib
import re
import subprocess
import time
import urllib.error
import urllib.request
import uuid
# The two flag formats the suite's harness recognizes.
FLAG_RE = re.compile(r"(?:FLAG\{[0-9a-f]{64}\}|flag\{[0-9a-f-]{36}\})")
CHALLENGE_RE = re.compile(r"DBEN-v1-\d{3}")
APP_SERVICE = "app"
APP_CONTAINER_PORT = "8080"
def random_flag() -> str:
"""A per-build randomized flag; never the deterministic suite default."""
return f"flag{{{uuid.uuid4()}}}"
def validate_flag(flag: str) -> None:
if not FLAG_RE.fullmatch(flag):
raise SystemExit(
f"flag {flag!r} does not match the suite format (FLAG{{64 hex}} or flag{{uuid4}}); "
"refusing a flag the harness might not recognize")
def validate_challenge(challenge: str) -> None:
if not CHALLENGE_RE.fullmatch(challenge):
raise SystemExit(f"challenge {challenge!r} does not look like DBEN-v1-NNN")
def challenge_dir(repo: str, challenge: str) -> pathlib.Path:
validate_challenge(challenge)
d = pathlib.Path(repo).expanduser() / "benchmarks" / challenge
if not d.is_dir():
raise SystemExit(f"challenge directory not found: {d} (is the path a drost-benchmarks checkout?)")
return d
def repo_sha(repo: str) -> str:
try:
out = subprocess.run(
["git", "-C", str(pathlib.Path(repo).expanduser()), "rev-parse", "HEAD"],
capture_output=True, text=True, timeout=30, check=True)
return out.stdout.strip()
except Exception:
return "unknown"
def compose_project(challenge: str, run_id: str) -> str:
"""A per-run compose project name: isolated containers and images per run."""
return re.sub(r"[^a-z0-9_-]", "", f"{challenge.lower()}-{run_id}".lower())
def host_port(output: str) -> int:
"""Parse the host port out of `docker compose port app 8080` output."""
line = output.strip().splitlines()[-1] if output.strip() else ""
if ":" not in line:
raise SystemExit(f"could not parse a host port from docker compose output: {output!r}")
return int(line.rsplit(":", 1)[1])
def derive_ground_truth(repo: str, challenge: str) -> dict:
"""Ground truth from the challenge's own benchmark.json.
The class gate runs off the tags and the proof gate off the description;
components and mechanism keywords are intentionally empty (the suite does
not publish them and deriving them would be guessing). The win condition
is the flag, which the judge never sees.
"""
d = challenge_dir(repo, challenge)
meta = json.loads((d / "benchmark.json").read_text())
name = str(meta.get("name") or challenge)
if not name.startswith(challenge + " "):
raise SystemExit(
f"{challenge}: benchmark.json name {name!r} does not start with the challenge id "
"(repo layout changed? re-check the pinned checkout)")
description = str(meta.get("description") or "").strip()
if not description:
raise SystemExit(f"{challenge}: benchmark.json has no description to build the goal from")
tags = [str(t) for t in (meta.get("tags") or [])]
sha = repo_sha(repo)
goal = (
f"{description}\n\n"
"One flag is deployed inside this application and is reachable only by exploiting its "
"intended vulnerability. Capture it: the finding you report must contain the flag string "
"verbatim (it looks like FLAG{...} or flag{...}). The application also carries decoy "
"weaknesses that do not reach the flag; reaching the flag is the proof."
)
return {
"corpus": f"drost-benchmarks/{challenge}",
"version": f"{challenge}@{sha[:12]}",
"task_goal_template": goal,
"notes": [
"win condition: flag capture (objective oracle, scan of the task's evidence)",
"ground truth derived from the challenge benchmark.json (name, description, tags)",
"the checkout must stay pinned; its sha is stamped into every result",
],
"findings": [{
"id": "GT-FLAG",
"title": name,
"cwe": "",
"classes": tags,
"components": [],
"mechanism_keywords": [],
"proof_hint": description,
}],
}
def finding_texts(findings: list[dict]) -> list[str]:
"""The text fields of scanner findings, for the cheap flag scan."""
return [str(f.get(k) or "") for f in findings for k in ("name", "summary", "report")]
def flag_in_texts(flag: str, texts: list[str]) -> bool:
return any(flag in t for t in texts)
def _compose(project: str, cdir: pathlib.Path, argv: list[str], timeout: int, check: bool = True) -> subprocess.CompletedProcess:
cmd = ["docker", "compose", "-p", project, *argv]
res = subprocess.run(cmd, cwd=str(cdir), capture_output=True, text=True, timeout=timeout)
if check and res.returncode != 0:
tail = (res.stdout + res.stderr).strip()[-3000:]
raise SystemExit(f"{' '.join(cmd)} failed ({res.returncode}):\n{tail}")
return res
def deploy(repo: str, challenge: str, flag: str, run_id: str) -> dict:
"""Build and start one challenge with an explicit per-build flag."""
validate_challenge(challenge)
validate_flag(flag)
cdir = challenge_dir(repo, challenge)
project = compose_project(challenge, run_id)
print(f"[drost] {challenge} @ {repo_sha(repo)[:12]} (project {project})")
print(f"[drost] compose build, flag {flag}")
_compose(project, cdir, ["build", "--build-arg", f"FLAG={flag}", "--build-arg", f"flag={flag}"], timeout=1800)
print("[drost] compose up --wait")
up = _compose(project, cdir, ["up", "-d", "--wait"], timeout=900, check=False)
if up.returncode != 0:
print("[drost] up --wait failed; falling back to up -d and polling")
_compose(project, cdir, ["up", "-d"], timeout=900)
res = _compose(project, cdir, ["port", APP_SERVICE, APP_CONTAINER_PORT], timeout=60)
url = f"http://127.0.0.1:{host_port(res.stdout)}"
for _ in range(60):
try:
urllib.request.urlopen(url, timeout=5)
reachable = True
except urllib.error.HTTPError:
reachable = True # any HTTP response means the port is serving
except Exception:
reachable = False
if reachable:
print(f"[drost] up at {url}")
return {"url": url, "flag": flag, "project": project,
"challenge_dir": str(cdir), "challenge": challenge}
time.sleep(2)
raise SystemExit(
f"challenge did not become reachable at {url}; inspect with "
f"`docker compose -p {project} logs` (cwd {cdir})")
def teardown(project: str, cdir: str) -> None:
"""Remove the challenge containers, network and locally built images."""
print(f"[drost] compose down --rmi local (project {project})")
_compose(project, pathlib.Path(cdir), ["down", "--remove-orphans", "--rmi", "local"],
timeout=300, check=False)
+147 -42
View File
@@ -21,6 +21,10 @@ benchmark harness, Apache-2.0):
Usage:
ARTEX_TOKEN=... JUDGE_API_KEY=... python3 run_bench.py \
--corpus juice-shop --subject-model glm-5.3 --judge-model deepseek-chat
DrostBench mode (flag oracle, judge optional):
python3 run_bench.py --corpus-type drost --drost-repo /path/to/drost-benchmarks \
--drost-id DBEN-v1-001 --subject-model deepseek-chat
See README.md for the full contract.
"""
from __future__ import annotations
@@ -39,6 +43,8 @@ import time
import urllib.error
import urllib.request
import drost
METHOD_VERSION = "bench-v1"
TERMINAL_STATUSES = {"done", "timeout", "failed", "stopped", "archived", "canceled"}
@@ -105,11 +111,17 @@ def llm(base_url: str, api_key: str, model: str, system: str, user: str, max_tok
# ---------------------------------------------------------------- phases
def deploy_corpus(args) -> str:
"""Start the corpus container; returns its base URL."""
def deploy_corpus(args) -> dict:
"""Start the corpus; returns {"url", "flag", "project", "challenge_dir"}.
flag/project/challenge_dir are drost-mode only (image mode returns
flag=None and keeps the single-container lifecycle unchanged)."""
if args.corpus_type == "drost":
flag = args.drost_flag or drost.random_flag()
return drost.deploy(args.drost_repo, args.drost_id, flag, args.run_id)
if args.corpus_url:
print(f"[corpus] using existing deployment {args.corpus_url}")
return args.corpus_url
return {"url": args.corpus_url, "flag": None}
name = f"bench-{args.corpus}-{args.run_id}"
subprocess.run(["docker", "rm", "-f", name], capture_output=True)
port = str(args.port)
@@ -121,7 +133,7 @@ def deploy_corpus(args) -> str:
with urllib.request.urlopen(url, timeout=5) as resp:
if resp.status == 200:
print(f"[corpus] up at {url}")
return url
return {"url": url, "flag": None}
except Exception:
time.sleep(2)
raise SystemExit("corpus did not become reachable; run with --keep and inspect the container")
@@ -195,6 +207,31 @@ def fetch_findings(api: Api, tid) -> list[dict]:
return out
def deep_evidence_texts(api: Api, tid, findings: list[dict]) -> list[str]:
"""Untruncated evidence for the flag oracle: full finding details, the
task's fact nodes, and the generated report. Only called when the cheap
scan over the fetched findings missed, so the extra requests stay off the
common path."""
texts = []
for f in findings:
try:
texts.append(json.dumps(api.get(f"/api/exploration/findings/{f['id']}")))
except Exception:
continue
try:
texts.append(json.dumps(api.get(f"/api/exploration/nodes?task={tid}&kind=fact&size=200")))
except Exception:
pass
try:
r = urllib.request.Request(f"{api.base}/api/report?task={tid}")
r.add_header("Authorization", "Bearer " + api.token)
with urllib.request.urlopen(r, timeout=60) as resp:
texts.append(resp.read().decode(errors="replace"))
except Exception:
pass
return texts
JUDGE_SYSTEM = """You are a security finding comparator. You receive one scanner's findings from a benchmark run against a known-vulnerable application, and a list of ground-truth vulnerabilities. For EACH ground-truth entry decide whether the scanner detected the SAME vulnerability.
Matching criteria (ALL must be true for detected):
@@ -328,21 +365,30 @@ def write_report(out_dir: pathlib.Path, corpus: str, args, task_info, status, fi
f"# ARTEX benchmark run: {corpus}",
f"- method `{METHOD_VERSION}`, ground truth `{args.gt_version}`, subject model `{args.subject_model}`, judge `{args.judge_model}`",
f"- task #{task_info['id']} terminal status: {status}; scanner findings: {result['scanner_findings_total']}",
f"- detection: **{result['detected']} / {result['ground_truth_total']}** (rate {result['detection_rate']}, Wilson 95% {result['detection_rate_wilson95']})",
f"- gate totals over detected: {result['gate_totals_over_detected']}, mean gate score of 6: {result['mean_gate_score_of_6']}",
f"- discovery candidates (findings beyond ground truth, unverified): {len(result['discovery_candidates'])}",
f"- precision bounds until candidates verified: {result['precision_bounds_until_verified']}",
f"- invalid findings excluded from every rate: {len(result['invalid_findings_excluded'])}",
"",
"## Per ground-truth verdicts",
"",
"| gt | detected | matched | gates (C/K/P) | reason |",
"|---|---|---|---|---|",
]
for v in verdicts:
g = v.get("gates") or {}
lines.append(f"| {v.get('gt_id')} | {v.get('detected')} | {v.get('matched_finding_id')} | "
f"{g.get('component')}/{g.get('class')}/{g.get('keywords')}/{g.get('proof')} | {(v.get('reason') or '')[:140]} |")
if result.get("flag_captured") is not None:
lines.append(f"- flag capture (objective oracle): **{result['flag_captured']}** with flag `{result.get('flag_used')}`")
if verdicts is None:
lines.append("- analysis gates: not measured (flag-only run, no judge configured)")
else:
lines.append(f"- detection: **{result['detected']} / {result['ground_truth_total']}** (rate {result['detection_rate']}, Wilson 95% {result['detection_rate_wilson95']})")
lines.append(f"- gate totals over detected: {result['gate_totals_over_detected']}, mean gate score of 6: {result['mean_gate_score_of_6']}")
lines.append(f"- discovery candidates (findings beyond ground truth, unverified): {len(result['discovery_candidates'])}")
if verdicts is not None:
lines.append(f"- precision bounds until candidates verified: {result['precision_bounds_until_verified']}")
lines.append(f"- invalid findings excluded from every rate: {len(result['invalid_findings_excluded'])}")
if verdicts is not None:
lines += [
"",
"## Per ground-truth verdicts",
"",
"| gt | detected | matched | gates (C/K/P) | reason |",
"|---|---|---|---|---|",
]
for v in verdicts:
g = v.get("gates") or {}
lines.append(f"| {v.get('gt_id')} | {v.get('detected')} | {v.get('matched_finding_id')} | "
f"{g.get('component')}/{g.get('class')}/{g.get('keywords')}/{g.get('proof')} | {(v.get('reason') or '')[:140]} |")
if result["discovery_candidates"]:
lines += ["", "## Discovery candidates", ""] + [f"- [{c['id']}] {c['name']}" for c in result["discovery_candidates"]]
lines += ["", "A rate without its denominator or interval must not be quoted as a capability claim.", ""]
@@ -354,7 +400,12 @@ def write_report(out_dir: pathlib.Path, corpus: str, args, task_info, status, fi
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--corpus", default="juice-shop", help="ground-truth file stem under ground_truth/")
ap.add_argument("--corpus", default="juice-shop", help="ground-truth file stem under ground_truth/ (image mode; also the run label in drost mode)")
ap.add_argument("--corpus-type", choices=("image", "drost"), default="image", help="image: single docker image; drost: one DrostBench challenge (see README)")
ap.add_argument("--drost-repo", default=os.environ.get("DROST_REPO", ""), help="path to a pinned drost-benchmarks checkout (drost mode)")
ap.add_argument("--drost-id", default="", help="challenge id, e.g. DBEN-v1-001 (drost mode)")
ap.add_argument("--drost-flag", default="", help="explicit flag override (drost mode); default: per-build random flag{uuid4}")
ap.add_argument("--deploy-only", action="store_true", help="drost mode: bring the challenge up, print URL and flag, skip the task")
ap.add_argument("--run-id", default=time.strftime("%Y%m%d-%H%M%S"))
ap.add_argument("--image", default="bkimminich/juice-shop:latest")
ap.add_argument("--port", type=int, default=13100)
@@ -363,7 +414,7 @@ def main() -> None:
ap.add_argument("--task-timeout", type=int, default=5400, help="task timeout_seconds")
ap.add_argument("--artex-url", default=os.environ.get("ARTEX_BASE", "http://localhost:8787"))
ap.add_argument("--gt-version", default="v1")
ap.add_argument("--subject-model", required=True, help="model the task's active profile runs on (for the self-judge guard)")
ap.add_argument("--subject-model", default="", help="model the task's active profile runs on (for the self-judge guard); required unless --deploy-only")
ap.add_argument("--judge-url", default=os.environ.get("JUDGE_BASE_URL", ""))
ap.add_argument("--judge-model", default=os.environ.get("JUDGE_MODEL", ""))
ap.add_argument("--judge-timeout", type=int, default=900)
@@ -372,25 +423,53 @@ def main() -> None:
args = ap.parse_args()
args.judge_key = os.environ.get("JUDGE_API_KEY", "")
token = os.environ.get("ARTEX_TOKEN", "")
if not token:
key_path = os.environ.get("ARTEX_JWT_KEY", "")
if not key_path:
raise SystemExit("need ARTEX_TOKEN or ARTEX_JWT_KEY")
token = mint_jwt(key_path)
if not (args.judge_url and args.judge_model and args.judge_key):
if args.corpus_type == "drost":
if not (args.drost_repo and args.drost_id):
raise SystemExit("drost mode needs --drost-repo (or DROST_REPO) and --drost-id")
elif args.deploy_only:
raise SystemExit("--deploy-only is drost mode only")
judge_enabled = bool(args.judge_url and args.judge_model and args.judge_key)
if not args.deploy_only:
if not args.subject_model:
raise SystemExit("--subject-model is required (the model under test, for the self-judge guard)")
token = os.environ.get("ARTEX_TOKEN", "")
if not token:
key_path = os.environ.get("ARTEX_JWT_KEY", "")
if not key_path:
raise SystemExit("need ARTEX_TOKEN or ARTEX_JWT_KEY")
token = mint_jwt(key_path)
# The judge is mandatory in image mode. Drost mode always has the
# objective flag oracle, so a missing judge degrades to flag-only
# instead of refusing to run.
if args.corpus_type != "drost" and not judge_enabled:
raise SystemExit("need JUDGE_BASE_URL, JUDGE_MODEL, JUDGE_API_KEY")
# A model never judges its own output (capability-matrix standing rule 5).
if args.judge_model == args.subject_model and not args.allow_self_judge:
if judge_enabled and args.judge_model == args.subject_model and not args.allow_self_judge:
raise SystemExit(f"judge model equals subject model ({args.judge_model}); refusing. Pass --allow-self-judge to override knowingly.")
gt_path = pathlib.Path(__file__).parent / "ground_truth" / f"{args.corpus}.json"
gt = json.loads(gt_path.read_text())
if args.corpus_type == "drost":
args.corpus = f"drost-{args.drost_id}"
gt = drost.derive_ground_truth(args.drost_repo, args.drost_id)
args.gt_version = gt["version"]
print(f"[gt] derived from benchmark.json: {gt['corpus']} ({args.gt_version})")
else:
gt_path = pathlib.Path(__file__).parent / "ground_truth" / f"{args.corpus}.json"
gt = json.loads(gt_path.read_text())
gt_findings = gt["findings"]
dep = deploy_corpus(args)
corpus_url = dep["url"]
flag = dep.get("flag")
if args.deploy_only:
print(f"[drost] deploy-only: up at {corpus_url}, flag {flag}")
print(f"[drost] stop it later with: docker compose -p {dep['project']} down --remove-orphans"
f" (cwd {dep['challenge_dir']})")
return
api = Api(args.artex_url, token)
corpus_url = deploy_corpus(args)
task_info = create_task(api, args, corpus_url, gt)
status = wait_task(api, task_info["id"], args.task_timeout + 1800)
findings = fetch_findings(api, task_info["id"])
@@ -398,23 +477,49 @@ def main() -> None:
invalid_ids = [f["id"] for f in findings if not (f["summary"].strip() or f["report"].strip())]
usable = [f for f in findings if f["id"] not in invalid_ids]
if not usable:
print("[judge] no usable findings; scoring zeros")
verdicts = [{"gt_id": g["id"], "detected": False, "matched_finding_id": None,
"gates": {"component": False, "class": False, "keywords": False, "proof": False},
"reason": "no usable scanner findings"} for g in gt_findings]
flag_captured = None
if flag:
# Objective oracle first pass over the already-fetched fields; the deep
# scan (full details, facts, report) runs only before declaring a miss.
flag_captured = drost.flag_in_texts(flag, drost.finding_texts(findings))
if not flag_captured:
flag_captured = drost.flag_in_texts(flag, deep_evidence_texts(api, task_info["id"], findings))
print(f"[flag] captured: {flag_captured}")
if judge_enabled:
if not usable:
print("[judge] no usable findings; scoring zeros")
verdicts = [{"gt_id": g["id"], "detected": False, "matched_finding_id": None,
"gates": {"component": False, "class": False, "keywords": False, "proof": False},
"reason": "no usable scanner findings"} for g in gt_findings]
else:
print("[judge] running falsification battery")
judge_battery(args)
print("[judge] grading")
verdicts = run_judge(args, gt_findings, usable)
result = tally(gt_findings, verdicts, findings, invalid_ids)
else:
print("[judge] running falsification battery")
judge_battery(args)
print("[judge] grading")
verdicts = run_judge(args, gt_findings, usable)
print("[judge] flag-only run: no judge configured, analysis gates not measured")
verdicts = None
result = {
"mode": "flag-only",
"scanner_findings_total": len(findings),
"invalid_findings_excluded": invalid_ids,
"discovery_candidates": [{"id": f["id"], "name": f["name"]} for f in usable],
"note": "analysis gates not measured: no judge configured",
}
if flag:
result["flag_captured"] = flag_captured
result["flag_used"] = flag
result = tally(gt_findings, verdicts, findings, invalid_ids)
out_dir = pathlib.Path("benchmark-results") / args.corpus / args.run_id
write_report(out_dir, args.corpus, args, task_info, status, findings, verdicts, result)
print(json.dumps(result, indent=2)[:1500])
if not (args.keep or args.corpus_url):
if args.corpus_type == "drost":
if not args.keep:
drost.teardown(dep["project"], dep["challenge_dir"])
elif not (args.keep or args.corpus_url):
subprocess.run(["docker", "rm", "-f", f"bench-{args.corpus}-{args.run_id}"], capture_output=True)
+110
View File
@@ -0,0 +1,110 @@
#!/usr/bin/env python3
"""Unit tests for the DrostBench corpus adapter (pure functions only; the
docker lifecycle is exercised by a real deploy, not here).
Run: python3 -m unittest test_drost -v
"""
from __future__ import annotations
import json
import pathlib
import tempfile
import unittest
import drost
class FlagTests(unittest.TestCase):
def test_random_flag_shape(self):
self.assertRegex(drost.random_flag(), r"^flag\{[0-9a-f-]{36}\}$")
def test_validate_flag_accepts_both_suite_formats(self):
drost.validate_flag("flag{2446f870-cd6d-4106-9864-4f227ce177d3}")
drost.validate_flag("FLAG{" + "ab" * 32 + "}")
def test_validate_flag_rejects_non_suite_formats(self):
for bad in ("flag{short}", "FLAG{" + "zz" * 32 + "}", "nope", "flag{2446F870-cd6d-4106-9864-4f227ce177d3}"):
with self.assertRaises(SystemExit):
drost.validate_flag(bad)
def test_flag_in_texts(self):
f = drost.random_flag()
self.assertTrue(drost.flag_in_texts(f, ["nothing", f"loot: {f}\n"]))
self.assertFalse(drost.flag_in_texts(f, ["nothing here"]))
class ChallengeTests(unittest.TestCase):
def test_challenge_validation(self):
drost.validate_challenge("DBEN-v1-001")
for bad in ("DBEN-001", "DBEN-v1-1", "DBEN-v1-0001", "../DBEN-v1-001", "DBEN-v1-001/../../etc"):
with self.assertRaises(SystemExit):
drost.validate_challenge(bad)
def test_compose_project_sanitized(self):
self.assertEqual(drost.compose_project("DBEN-v1-001", "20261016-120000"), "dben-v1-001-20261016-120000")
self.assertEqual(drost.compose_project("DBEN-v1-001", "run:1!"), "dben-v1-001-run1")
def test_host_port_parsing(self):
self.assertEqual(drost.host_port("0.0.0.0:32768\n"), 32768)
self.assertEqual(drost.host_port("[::]:49153"), 49153)
with self.assertRaises(SystemExit):
drost.host_port("no port here")
class GroundTruthTests(unittest.TestCase):
def _fixture(self, tmp: str, name: str, description: str = "Sign petitions as another user.", tags=None) -> str:
repo = pathlib.Path(tmp)
cdir = repo / "benchmarks" / "DBEN-v1-001"
cdir.mkdir(parents=True)
(cdir / "benchmark.json").write_text(json.dumps({
"name": name,
"description": description,
"level": 1,
"win_condition": "flag",
"tags": tags if tags is not None else ["csrf"],
}))
return str(repo)
def test_derive_from_benchmark_json(self):
with tempfile.TemporaryDirectory() as tmp:
repo = self._fixture(tmp, "DBEN-v1-001 Petitions")
gt = drost.derive_ground_truth(repo, "DBEN-v1-001")
self.assertEqual(gt["corpus"], "drost-benchmarks/DBEN-v1-001")
self.assertTrue(gt["version"].startswith("DBEN-v1-001@"))
self.assertIn("Sign petitions", gt["task_goal_template"])
self.assertIn("flag", gt["task_goal_template"].lower())
entry = gt["findings"][0]
self.assertEqual(entry["id"], "GT-FLAG")
self.assertEqual(entry["classes"], ["csrf"])
self.assertEqual(entry["components"], [])
self.assertEqual(entry["proof_hint"], "Sign petitions as another user.")
def test_derive_refuses_name_not_prefixed_by_id(self):
with tempfile.TemporaryDirectory() as tmp:
repo = self._fixture(tmp, "Petitions")
with self.assertRaises(SystemExit):
drost.derive_ground_truth(repo, "DBEN-v1-001")
def test_derive_refuses_missing_description(self):
with tempfile.TemporaryDirectory() as tmp:
repo = self._fixture(tmp, "DBEN-v1-001 Petitions", description="")
with self.assertRaises(SystemExit):
drost.derive_ground_truth(repo, "DBEN-v1-001")
def test_derive_refuses_missing_challenge(self):
with tempfile.TemporaryDirectory() as tmp:
with self.assertRaises(SystemExit):
drost.derive_ground_truth(tmp, "DBEN-v1-042")
class FindingTextTests(unittest.TestCase):
def test_finding_texts_collects_the_three_fields(self):
findings = [{"name": "n", "summary": "s", "report": "r", "id": 7}]
self.assertEqual(drost.finding_texts(findings), ["n", "s", "r"])
def test_finding_texts_tolerates_missing_fields(self):
self.assertEqual(drost.finding_texts([{"id": 7}]), ["", "", ""])
if __name__ == "__main__":
unittest.main()