iohex/cybergym
1
0
Fork 0
mirror of https://github.com/sunblaze-ucb/cybergym.git synced 2026-10-12 08:42:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: binary only mode

Browse Source
This commit is contained in:
stneng
2026-02-09 16:18:46 -08:00
parent f86dfd7ca9
commit bcbe562ee6
9 changed files with 109 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
README.md
View File

@@ -27,8 +27,6 @@ Download the PoC submission server data
1. Full data
```bash
python scripts/server_data/download.py --tasks-file ./cybergym_data/tasks.json
bash scripts/server_data/download_chunks.sh
7z x cybergym-oss-fuzz-data.7z
```
2. Subset data
@@ -49,8 +47,13 @@ oss-fuzz:385167047
Download the subset data
```bash
python scripts/server_data/download_subset.py
wget https://huggingface.co/datasets/sunblaze-ucb/cybergym-server/resolve/main/cybergym-oss-fuzz-data-subset.7z
7z x cybergym-oss-fuzz-data-subset.7z
```
### Download Server Data (binary only mode)
```
python scripts/server_data/download_binary_only_runners.py
wget https://huggingface.co/datasets/sunblaze-ucb/cybergym-server-binary/resolve/main/cybergym-server-data.7z
7z x cybergym-server-data.7z
```
## Evaluation
@@ -58,11 +61,20 @@ Start the PoC submission server:
```bash
PORT=8666 # port of the server
POC_SAVE_DIR=./server_poc # dir to save the pocs
CYBERGYM_SERVER_DATA_DIR=./oss-fuzz-data
python3 -m cybergym.server \
--host 0.0.0.0 --port $PORT \
--log_dir $POC_SAVE_DIR --db_path $POC_SAVE_DIR/poc.db
```
Start the PoC submission server (binary only mode):
```bash
PORT=8666 # port of the server
POC_SAVE_DIR=./server_poc # dir to save the pocs
CYBERGYM_SERVER_DATA_DIR=./cybergym-server-data
python3 -m cybergym.server \
--host 0.0.0.0 --port $PORT \
--log_dir $POC_SAVE_DIR --db_path $POC_SAVE_DIR/poc.db \
--cybergym_oss_fuzz_path $CYBERGYM_SERVER_DATA_DIR
--binary_dir $CYBERGYM_SERVER_DATA_DIR
```
Test:

2
pyproject.toml
View File

@@ -1,6 +1,6 @@
[project]
name = "cybergym"
version = "0.1.0"
version = "0.2.0"
description = "A Python project for cyber security training"
readme = "README.md"
requires-python = ">=3.12"

4
scripts/server_data/binary_only_runners.txt
View File

@@ -1,4 +0,0 @@
cybergym/oss-fuzz-base-runner:latest
cybergym/oss-fuzz-base-runner:20200102
cybergym/oss-fuzz-base-runner:20190802
cybergym/oss-fuzz-base-runner:20220102

17
scripts/server_data/download.py
View File

@@ -46,11 +46,18 @@ if __name__ == "__main__":
with open(args.tasks_file) as f:
tasks = json.load(f)
tags = []
tags_arvo = []
tags_ossfuzz = []
for task in tasks:
if task["task_id"].split(":")[0] == "arvo":
arvo_id = task["task_id"].split(":")[-1]
tags.append(f"{arvo_id}-vul")
tags.append(f"{arvo_id}-fix")
if tags:
pull_images("n132/arvo", tags, args.max_workers)
tags_arvo.append(f"{arvo_id}-vul")
tags_arvo.append(f"{arvo_id}-fix")
if task["task_id"].split(":")[0] == "ossfuzz":
ossfuzz_id = task["task_id"].split(":")[-1]
tags_ossfuzz.append(f"{ossfuzz_id}-vul")
tags_ossfuzz.append(f"{ossfuzz_id}-fix")
if tags_arvo:
pull_images("n132/arvo", tags_arvo, args.max_workers)
if tags_ossfuzz:
pull_images("cybergym/oss-fuzz", tags_ossfuzz, args.max_workers)

36
scripts/server_data/download_binary_only_runners.py Normal file
View File

@@ -0,0 +1,36 @@
import argparse
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
import docker
import json
client = docker.from_env()
def pull_images(repo, tags, max_workers=1):
def _pull(tag):
image = f"{repo}:{tag}"
print(f"Pulling {image}...")
try:
client.images.pull(repo, tag=tag)
print(f"Successfully pulled {image}")
except docker.errors.APIError as e:
print(f"Failed to pull {image}: {e}")
if max_workers == 1:
for tag in tags:
_pull(tag)
return
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {executor.submit(_pull, tag): tag for tag in tags}
for future in as_completed(futures):
try:
future.result()
except Exception as e:
tag = futures[future]
print(f"Unexpected error pulling {repo}:{tag}: {e}")
if __name__ == "__main__":
pull_images("cybergym/oss-fuzz-base-runner", ["latest", "20200102", "20190802", "20220102"], 1)

17
scripts/server_data/download_subset.py
View File

@@ -54,11 +54,18 @@ if __name__ == "__main__":
pull_images("cybergym/oss-fuzz-base-runner", ["latest"], 1)
tags = []
tags_arvo = []
tags_ossfuzz = []
for task_id in task_ids:
if task_id.split(":")[0] == "arvo":
arvo_id = task_id.split(":")[-1]
tags.append(f"{arvo_id}-vul")
tags.append(f"{arvo_id}-fix")
if tags:
pull_images("n132/arvo", tags, args.max_workers)
tags_arvo.append(f"{arvo_id}-vul")
tags_arvo.append(f"{arvo_id}-fix")
if task_id.split(":")[0] == "oss-fuzz":
ossfuzz_id = task_id.split(":")[-1]
tags_ossfuzz.append(f"{ossfuzz_id}-vul")
tags_ossfuzz.append(f"{ossfuzz_id}-fix")
if tags_arvo:
pull_images("n132/arvo", tags_arvo, args.max_workers)
if tags_ossfuzz:
pull_images("cybergym/oss-fuzz", tags_ossfuzz, args.max_workers)

2
scripts/verify_agent_result.py
View File

@@ -11,6 +11,7 @@ API_KEY = os.getenv("CYBERGYM_API_KEY")
API_KEY_NAME = "X-API-Key"
logger = logging.getLogger(__name__)
def run_verify(agent_id: str, server: str):
with httpx.Client(base_url=server, timeout=1200) as client:
headers = {
@@ -28,6 +29,7 @@ def run_verify(agent_id: str, server: str):
except Exception as e:
logger.error(f"Error during verification for agent {agent_id}: {e}")
def load_results(pocdb_path: Path, agent_id: str):
engine = init_engine(pocdb_path)
with Session(engine) as session:

14
src/cybergym/server/__main__.py
View File

@@ -75,8 +75,10 @@ def submit_vul(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated
except Exception:
raise HTTPException(status_code=400, detail="Invalid metadata format") from None
payload.data = file_content
use2 = bool(server_conf.binary_dir)
res = submit_poc(db, payload, mode="vul", log_dir=server_conf.log_dir, salt=server_conf.salt, use2=use2)
binary_only_mode = bool(server_conf.binary_dir)
res = submit_poc(
db, payload, mode="vul", log_dir=server_conf.log_dir, salt=server_conf.salt, binary_only_mode=binary_only_mode
)
res = _post_process_result(res, payload.require_flag)
return res
@@ -96,8 +98,10 @@ def submit_fix(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated
except Exception:
raise HTTPException(status_code=400, detail="Invalid metadata format") from None
payload.data = file_content
use2 = bool(server_conf.binary_dir)
res = submit_poc(db, payload, mode="fix", log_dir=server_conf.log_dir, salt=server_conf.salt, use2=use2)
binary_only_mode = bool(server_conf.binary_dir)
res = submit_poc(
db, payload, mode="fix", log_dir=server_conf.log_dir, salt=server_conf.salt, binary_only_mode=binary_only_mode
)
res = _post_process_result(res, payload.require_flag)
return res
@@ -120,7 +124,7 @@ def verify_all_pocs_for_agent_id(db: SessionDep, query: VerifyPocs):
raise HTTPException(status_code=404, detail="No records found for this agent_id")
for record in records:
run_poc_id(db, server_conf.log_dir, record.poc_id)
run_poc_id(db, server_conf.log_dir, record.poc_id, binary_only_mode=bool(server_conf.binary_dir))
return {
"message": f"All {len(records)} PoCs for this agent_id have been verified",

30
src/cybergym/server/server_utils.py
View File

@@ -104,7 +104,7 @@ def run_container(
return exit_code, docker_output
def run_container2(
def run_container_binary(
task_id: str,
poc_path: Path,
mode: Literal["vul", "fix"],
@@ -120,7 +120,7 @@ def run_container2(
container = None
if subset == "arvo":
runner_image_file = data_dir / "arvo" / subid / "runner_image.txt"
runner_image_file = data_dir / "arvo" / subid / mode / "runner"
if runner_image_file.exists():
runner_image = runner_image_file.read_text().strip()
bin_dir = data_dir / "arvo" / subid / mode
@@ -148,8 +148,8 @@ def run_container2(
if not is_integer(subid):
raise HTTPException(status_code=400, detail="Invalid task_id format for oss-fuzz")
oss_fuzz_path = data_dir / "oss-fuzz"
out_dir = oss_fuzz_path / f"{subid}-{mode}/out"
meta_file = oss_fuzz_path / f"{subid}-{mode}/metadata.json"
out_dir = oss_fuzz_path / subid / mode / "out"
meta_file = oss_fuzz_path / subid / mode / "metadata.json"
with open(meta_file) as f:
metadata = json.load(f)
fuzzer_name = metadata["fuzz_target"]
@@ -194,7 +194,7 @@ def get_poc_storage_path(poc_id: str, log_dir: Path):
return log_dir / poc_id[:2] / poc_id[2:4] / poc_id
def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str, use2: bool = False):
def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str, binary_only_mode: bool = False):
# TODO: limit output size for return
if not verify_task(payload.task_id, payload.agent_id, payload.checksum, salt=salt):
raise HTTPException(status_code=400, detail="Invalid checksum")
@@ -249,8 +249,10 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
)
# Run the PoC
if use2:
exit_code, docker_output = run_container2(payload.task_id, poc_bin_file, mode, data_dir=server_conf.binary_dir)
if binary_only_mode:
exit_code, docker_output = run_container_binary(
payload.task_id, poc_bin_file, mode, data_dir=server_conf.binary_dir
)
else:
exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode)
output_file = poc_dir / f"output.{mode}"
@@ -268,7 +270,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
return res
def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, use2: bool = False):
def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, binary_only_mode: bool = False):
records = db.query(PoCRecord).filter_by(poc_id=poc_id).all()
if len(records) != 1:
raise HTTPException(status_code=500, detail=f"{len(records)} PoC records for same poc_id found")
@@ -281,8 +283,10 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, use
if rerun or record.vul_exit_code is None:
# Run the PoC
if use2:
exit_code, docker_output = run_container2(record.task_id, poc_path, "vul", data_dir=server_conf.binary_dir)
if binary_only_mode:
exit_code, docker_output = run_container_binary(
record.task_id, poc_path, "vul", data_dir=server_conf.binary_dir
)
else:
exit_code, docker_output = run_container(record.task_id, poc_path, "vul")
with open(poc_dir / "output.vul", "wb") as f:
@@ -295,8 +299,10 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, use
if rerun or record.fix_exit_code is None:
# Run the PoC
if use2:
exit_code, docker_output = run_container2(record.task_id, poc_path, "fix", data_dir=server_conf.binary_dir)
if binary_only_mode:
exit_code, docker_output = run_container_binary(
record.task_id, poc_path, "fix", data_dir=server_conf.binary_dir
)
else:
exit_code, docker_output = run_container(record.task_id, poc_path, "fix")
with open(poc_dir / "output.fix", "wb") as f: