feat: binary only mode

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:

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"

scripts/server_data/binary_only_runners.txt

@@ -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

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_arvo.append(f"{arvo_id}-vul")
-            tags.append(f"{arvo_id}-fix")
+            tags_arvo.append(f"{arvo_id}-fix")
-    if tags:
+        if task["task_id"].split(":")[0] == "ossfuzz":
-        pull_images("n132/arvo", tags, args.max_workers)
+            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)

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)

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_arvo.append(f"{arvo_id}-vul")
-            tags.append(f"{arvo_id}-fix")
+            tags_arvo.append(f"{arvo_id}-fix")
-    if tags:
+        if task_id.split(":")[0] == "oss-fuzz":
-        pull_images("n132/arvo", tags, args.max_workers)
+            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)

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)
+    load_results(args.pocdb_path, args.agent_id)

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)
+    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, use2=use2)
+    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)
+    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, use2=use2)
+    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",

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"
+        out_dir = oss_fuzz_path / subid / mode / "out"
-        meta_file = oss_fuzz_path / f"{subid}-{mode}/metadata.json"
+        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:
+    if binary_only_mode:
-        exit_code, docker_output = run_container2(payload.task_id, poc_bin_file, mode, data_dir=server_conf.binary_dir)
+        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:
+        if binary_only_mode:
-            exit_code, docker_output = run_container2(record.task_id, poc_path, "vul", data_dir=server_conf.binary_dir)
+            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:
+        if binary_only_mode:
-            exit_code, docker_output = run_container2(record.task_id, poc_path, "fix", data_dir=server_conf.binary_dir)
+            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: