From bcbe562ee6247a8bdc64e4cba11167b6fcd66603 Mon Sep 17 00:00:00 2001 From: stneng Date: Mon, 9 Feb 2026 16:18:46 -0800 Subject: [PATCH] feat: binary only mode --- README.md | 24 +++++++++---- pyproject.toml | 2 +- scripts/server_data/binary_only_runners.txt | 4 --- scripts/server_data/download.py | 17 ++++++--- .../download_binary_only_runners.py | 36 +++++++++++++++++++ scripts/server_data/download_subset.py | 17 ++++++--- scripts/verify_agent_result.py | 4 ++- src/cybergym/server/__main__.py | 14 +++++--- src/cybergym/server/server_utils.py | 30 +++++++++------- 9 files changed, 109 insertions(+), 39 deletions(-) delete mode 100644 scripts/server_data/binary_only_runners.txt create mode 100644 scripts/server_data/download_binary_only_runners.py diff --git a/README.md b/README.md index 33a24bd..6745f50 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/pyproject.toml b/pyproject.toml index c9712ae..bd5bd73 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/scripts/server_data/binary_only_runners.txt b/scripts/server_data/binary_only_runners.txt deleted file mode 100644 index b43c37f..0000000 --- a/scripts/server_data/binary_only_runners.txt +++ /dev/null @@ -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 diff --git a/scripts/server_data/download.py b/scripts/server_data/download.py index 1d90a7f..769629c 100644 --- a/scripts/server_data/download.py +++ b/scripts/server_data/download.py @@ -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) diff --git a/scripts/server_data/download_binary_only_runners.py b/scripts/server_data/download_binary_only_runners.py new file mode 100644 index 0000000..e22d309 --- /dev/null +++ b/scripts/server_data/download_binary_only_runners.py @@ -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) diff --git a/scripts/server_data/download_subset.py b/scripts/server_data/download_subset.py index 6b8572c..fd45151 100644 --- a/scripts/server_data/download_subset.py +++ b/scripts/server_data/download_subset.py @@ -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) diff --git a/scripts/verify_agent_result.py b/scripts/verify_agent_result.py index 589be45..f76c7fe 100644 --- a/scripts/verify_agent_result.py +++ b/scripts/verify_agent_result.py @@ -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: @@ -60,4 +62,4 @@ if __name__ == "__main__": args = parser.parse_args() run_verify(args.agent_id, args.server) - load_results(args.pocdb_path, args.agent_id) \ No newline at end of file + load_results(args.pocdb_path, args.agent_id) diff --git a/src/cybergym/server/__main__.py b/src/cybergym/server/__main__.py index b079bb5..8dae39f 100644 --- a/src/cybergym/server/__main__.py +++ b/src/cybergym/server/__main__.py @@ -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", diff --git a/src/cybergym/server/server_utils.py b/src/cybergym/server/server_utils.py index 289e301..d39376e 100644 --- a/src/cybergym/server/server_utils.py +++ b/src/cybergym/server/server_utils.py @@ -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: