From 73632517cd769da73a062459d056b82ec7bc6a54 Mon Sep 17 00:00:00 2001 From: zhun Date: Tue, 13 Jan 2026 17:14:45 +0000 Subject: [PATCH 1/5] update cybergym/oss-fuzz subset --- src/cybergym/server/__main__.py | 10 +- src/cybergym/server/server_utils.py | 151 +++++++--------------------- 2 files changed, 40 insertions(+), 121 deletions(-) diff --git a/src/cybergym/server/__main__.py b/src/cybergym/server/__main__.py index d171dd6..63509c1 100644 --- a/src/cybergym/server/__main__.py +++ b/src/cybergym/server/__main__.py @@ -17,7 +17,6 @@ from cybergym.task.types import DEFAULT_SALT SALT = DEFAULT_SALT LOG_DIR = Path("./logs") DB_PATH = Path("./poc.db") -OSS_FUZZ_PATH = Path("./oss-fuzz-data") API_KEY = "cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d" API_KEY_NAME = "X-API-Key" @@ -66,7 +65,7 @@ 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.file.read() - res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT, oss_fuzz_path=OSS_FUZZ_PATH) + res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT) res = _post_process_result(res, payload.require_flag) return res @@ -78,7 +77,7 @@ 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.file.read() - res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT, oss_fuzz_path=OSS_FUZZ_PATH) + res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT) res = _post_process_result(res, payload.require_flag) return res @@ -101,7 +100,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, LOG_DIR, record.poc_id, oss_fuzz_path=OSS_FUZZ_PATH) + run_poc_id(db, LOG_DIR, record.poc_id) return { "message": f"All {len(records)} PoCs for this agent_id have been verified", @@ -120,7 +119,6 @@ if __name__ == "__main__": parser.add_argument("--salt", type=str, default=SALT, help="Salt for checksum") parser.add_argument("--log_dir", type=Path, default=LOG_DIR, help="Directory to store logs") parser.add_argument("--db_path", type=Path, default=DB_PATH, help="Path to SQLite DB") - parser.add_argument("--cybergym_oss_fuzz_path", type=Path, default=OSS_FUZZ_PATH, help="Path to OSS-Fuzz") args = parser.parse_args() SALT = args.salt @@ -129,6 +127,4 @@ if __name__ == "__main__": DB_PATH = Path(args.db_path) - OSS_FUZZ_PATH = Path(args.cybergym_oss_fuzz_path) - uvicorn.run(app, host=args.host, port=args.port) diff --git a/src/cybergym/server/server_utils.py b/src/cybergym/server/server_utils.py index a9b0e79..153772b 100644 --- a/src/cybergym/server/server_utils.py +++ b/src/cybergym/server/server_utils.py @@ -1,6 +1,5 @@ import hashlib -import json -import os +import shlex from enum import IntEnum from pathlib import Path from typing import Literal @@ -40,19 +39,45 @@ def _post_process_result(res: dict, require_flag: bool = False): return res -def run_arvo_container( +def _image_and_command_from_task_id(task_id: str, mode: str) -> tuple[str, list[str]]: + if task_id.startswith("arvo:"): + arvo_id = get_arvo_id(task_id) + image = f"n132/arvo:{arvo_id}-{mode}" + command = ["/bin/arvo"] + elif task_id.startswith("oss-fuzz:"): + oss_fuzz_id = get_oss_fuzz_id(task_id) + image = f"cybergym/oss-fuzz:{oss_fuzz_id}-{mode}" + command = ["/usr/local/bin/run_poc"] + elif task_id.startswith("oss-fuzz-latest:"): + raise HTTPException(status_code=400, detail="oss-fuzz-latest does not support this operation") + else: + raise HTTPException(status_code=400, detail="Invalid task_id") + return image, command + + +def is_integer(s): + try: + int(s) + return True + except ValueError: + return False + + + +def run_container( + task_id: str, poc_path: Path, - arvo_id: str, mode: Literal["vul", "fix"], docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, cmd_timeout: int = DEFAULT_CMD_TIMEOUT, ): + image, cmd = _image_and_command_from_task_id(task_id, mode) + cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} {shlex.join(cmd)} 2>&1"] client = docker.from_env() container = None try: - cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} /bin/arvo 2>&1"] container = client.containers.run( - image=f"n132/arvo:{arvo_id}-{mode}", + image=image, command=cmd, volumes={str(poc_path.absolute()): {"bind": "/tmp/poc", "mode": "ro"}}, # noqa: S108 detach=True, @@ -74,109 +99,7 @@ def run_arvo_container( if container: container.remove(force=True) - return exit_code, docker_output - - -def is_integer(s): - try: - int(s) - return True - except ValueError: - return False - - -def run_oss_fuzz_container( - poc_path: Path, - oss_fuzz_id: str, - mode: Literal["vul", "fix"], - oss_fuzz_path: Path, - docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, - cmd_timeout: int = DEFAULT_CMD_TIMEOUT, -): - client = docker.from_env() - container = None - try: - if is_integer(oss_fuzz_id): - out_dir = Path(oss_fuzz_path, f"{oss_fuzz_id}-{mode}", "out") - else: - if mode == "fix": - raise HTTPException(status_code=400, detail="Fix mode is not supported for oss-fuzz-latest") - project, index = oss_fuzz_id.rsplit("-", 1) - out_dir = Path(oss_fuzz_path, project, "out") - volumes = {str(poc_path.absolute()): {"bind": "/testcase", "mode": "ro"}} - for filename in os.listdir(out_dir): - host_path = str(Path(out_dir, filename).absolute()) - container_path = os.path.join("/out", filename) - volumes[host_path] = {"bind": container_path, "mode": "ro"} - if is_integer(oss_fuzz_id): - meta_file = os.path.join(oss_fuzz_path, f"{oss_fuzz_id}-{mode}", "metadata.json") - with open(meta_file) as f: - metadata = json.load(f) - fuzzer_name = metadata["fuzz_target"] - else: - project, index = oss_fuzz_id.rsplit("-", 1) - meta_file = os.path.join(oss_fuzz_path, project, "metadata.json") - with open(meta_file) as f: - metadata = json.load(f) - fuzzer_name = metadata["fuzz_targets"][int(index)] - - cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} reproduce {fuzzer_name} 2>&1"] - container = client.containers.run( - image="cybergym/oss-fuzz-base-runner", - command=cmd, - volumes=volumes, - detach=True, - ) - out = container.logs(stdout=True, stderr=False, stream=True, follow=True) - exit_code = container.wait(timeout=docker_timeout)["StatusCode"] - if exit_code == 137: # Process killed by timeout - exit_code = CustomExitCode.Timeout - docker_output = b"" - else: - docker_output = b"".join(out) - except requests.exceptions.ReadTimeout: - raise HTTPException(status_code=500, detail="Timeout waiting for the program") from None - except DockerException as e: - raise HTTPException(status_code=500, detail=f"Running error: {e}") from None - except Exception as e: - raise HTTPException(status_code=500, detail=f"Unexpected error: {e}") from None - finally: - if container: - container.remove(force=True) - - return exit_code, docker_output - - -def run_container( - task_id: str, - poc_path: Path, - mode: Literal["vul", "fix"], - docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, - cmd_timeout: int = DEFAULT_CMD_TIMEOUT, - **kwargs, -): - if task_id.startswith("arvo:"): - arvo_id = get_arvo_id(task_id) - return run_arvo_container( - poc_path, - arvo_id, - mode, - docker_timeout=docker_timeout, - cmd_timeout=cmd_timeout, - ) - elif task_id.startswith("oss-fuzz:") or task_id.startswith("oss-fuzz-latest:"): - oss_fuzz_id = get_oss_fuzz_id(task_id) - oss_fuzz_path = kwargs.get("oss_fuzz_path") - return run_oss_fuzz_container( - poc_path, - oss_fuzz_id, - mode, - oss_fuzz_path, - docker_timeout=docker_timeout, - cmd_timeout=cmd_timeout, - ) - else: - raise HTTPException(status_code=400, detail="Invalid task_id") + return exit_code, docker_output def get_poc_storage_path(poc_id: str, log_dir: Path): @@ -184,7 +107,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, oss_fuzz_path: Path | None = None): +def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str): # 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") @@ -239,7 +162,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st ) # Run the PoC - exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode, oss_fuzz_path=oss_fuzz_path) + exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode) output_file = poc_dir / f"output.{mode}" with open(output_file, "wb") as f: f.write(docker_output) @@ -255,7 +178,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, oss_fuzz_path: Path | None = None): +def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: 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") @@ -268,7 +191,7 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, oss if rerun or record.vul_exit_code is None: # Run the PoC - exit_code, docker_output = run_container(record.task_id, poc_path, "vul", oss_fuzz_path=oss_fuzz_path) + exit_code, docker_output = run_container(record.task_id, poc_path, "vul") with open(poc_dir / "output.vul", "wb") as f: f.write(docker_output) update_poc_output(db, record, "vul", exit_code) @@ -279,7 +202,7 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, oss if rerun or record.fix_exit_code is None: # Run the PoC - exit_code, docker_output = run_container(record.task_id, poc_path, "fix", oss_fuzz_path=oss_fuzz_path) + exit_code, docker_output = run_container(record.task_id, poc_path, "fix") with open(poc_dir / "output.fix", "wb") as f: f.write(docker_output) update_poc_output(db, record, "fix", exit_code) From 18ea589df8035bd8023e80350a7c8b26b690feab Mon Sep 17 00:00:00 2001 From: zhun Date: Sun, 8 Feb 2026 09:42:54 -0800 Subject: [PATCH 2/5] add binary mode --- pyproject.toml | 1 + src/cybergym/server/__main__.py | 73 +++++++++++++++------- src/cybergym/server/server_utils.py | 95 +++++++++++++++++++++++++++-- src/cybergym/server/types.py | 26 +++++++- 4 files changed, 165 insertions(+), 30 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 946d4c1..c9712ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ server = [ "uvicorn>=0.34.0", "sqlalchemy>=2.0.0", "python-multipart>=0.0.20", + "pydantic-settings>=2.12.0", ] [project.urls] diff --git a/src/cybergym/server/__main__.py b/src/cybergym/server/__main__.py index 63509c1..b079bb5 100644 --- a/src/cybergym/server/__main__.py +++ b/src/cybergym/server/__main__.py @@ -11,14 +11,7 @@ from sqlalchemy.orm import Session from cybergym.server.pocdb import get_poc_by_hash, init_engine from cybergym.server.server_utils import _post_process_result, run_poc_id, submit_poc -from cybergym.server.types import Payload, PocQuery, VerifyPocs -from cybergym.task.types import DEFAULT_SALT - -SALT = DEFAULT_SALT -LOG_DIR = Path("./logs") -DB_PATH = Path("./poc.db") -API_KEY = "cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d" -API_KEY_NAME = "X-API-Key" +from cybergym.server.types import Payload, PocQuery, VerifyPocs, server_conf engine: Engine = None @@ -34,7 +27,7 @@ SessionDep = Annotated[Session, Depends(get_session)] @asynccontextmanager async def lifespan(app: FastAPI): global engine - engine = init_engine(DB_PATH) + engine = init_engine(server_conf.db_path) yield @@ -44,40 +37,67 @@ async def lifespan(app: FastAPI): app = FastAPI(lifespan=lifespan) -api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False) +api_key_header = APIKeyHeader(name=server_conf.api_key_name, auto_error=False) def get_api_key(api_key: str = Security(api_key_header)): - if api_key == API_KEY: + if api_key == server_conf.api_key: return api_key else: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not found") +def try_read_file(file: UploadFile, max_size_mb: int) -> bytes: + """Helper function to check file size.""" + max_size_bytes = max_size_mb * 1024 * 1024 + content = file.file.read(max_size_bytes + 1) + if len(content) > max_size_bytes: + raise HTTPException(status_code=413, detail=f"File too large. Maximum size allowed: {max_size_mb}MB") + return content + + public_router = APIRouter() private_router = APIRouter(dependencies=[Depends(get_api_key)]) @public_router.post("/submit-vul") def submit_vul(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated[UploadFile, File()]): + # Read and validate file size + try: + file_content = try_read_file(file, server_conf.max_file_size_mb) + except HTTPException: + raise + except Exception: + raise HTTPException(status_code=400, detail="Error reading file") from None + try: payload = Payload.model_validate_json(metadata) except Exception: raise HTTPException(status_code=400, detail="Invalid metadata format") from None - payload.data = file.file.read() - res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT) + 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) res = _post_process_result(res, payload.require_flag) return res @private_router.post("/submit-fix") def submit_fix(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated[UploadFile, File()]): + # Read and validate file size + try: + file_content = try_read_file(file, server_conf.max_file_size_mb) + except HTTPException: + raise + except Exception: + raise HTTPException(status_code=400, detail="Error reading file") from None + try: payload = Payload.model_validate_json(metadata) except Exception: raise HTTPException(status_code=400, detail="Invalid metadata format") from None - payload.data = file.file.read() - res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT) + 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) res = _post_process_result(res, payload.require_flag) return res @@ -100,7 +120,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, LOG_DIR, record.poc_id) + run_poc_id(db, server_conf.log_dir, record.poc_id) return { "message": f"All {len(records)} PoCs for this agent_id have been verified", @@ -116,15 +136,22 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="CyberGym Server") parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to run the server on") parser.add_argument("--port", type=int, default=8666, help="Port to run the server on") - parser.add_argument("--salt", type=str, default=SALT, help="Salt for checksum") - parser.add_argument("--log_dir", type=Path, default=LOG_DIR, help="Directory to store logs") - parser.add_argument("--db_path", type=Path, default=DB_PATH, help="Path to SQLite DB") + parser.add_argument("--salt", type=str, default=server_conf.salt, help="Salt for checksum") + parser.add_argument("--log_dir", type=Path, default=server_conf.log_dir, help="Directory to store logs") + parser.add_argument("--db_path", type=Path, default=server_conf.db_path, help="Path to SQLite DB") + parser.add_argument( + "--binary_dir", type=Path, default=server_conf.binary_dir, help="Directory to store target binaries" + ) + parser.add_argument( + "--max_file_size_mb", type=int, default=server_conf.max_file_size_mb, help="Maximum file size for uploads in MB" + ) args = parser.parse_args() - SALT = args.salt - LOG_DIR = args.log_dir - LOG_DIR.mkdir(parents=True, exist_ok=True) - DB_PATH = Path(args.db_path) + server_conf.salt = args.salt + server_conf.log_dir = args.log_dir + server_conf.log_dir.mkdir(parents=True, exist_ok=True) + server_conf.db_path = Path(args.db_path) + server_conf.binary_dir = args.binary_dir uvicorn.run(app, host=args.host, port=args.port) diff --git a/src/cybergym/server/server_utils.py b/src/cybergym/server/server_utils.py index 153772b..7e9288f 100644 --- a/src/cybergym/server/server_utils.py +++ b/src/cybergym/server/server_utils.py @@ -1,4 +1,6 @@ import hashlib +import json +import os import shlex from enum import IntEnum from pathlib import Path @@ -12,13 +14,14 @@ from fastapi import HTTPException from sqlalchemy.orm import Session from cybergym.server.pocdb import PoCRecord, get_or_create_poc, get_poc_by_hash, update_poc_output -from cybergym.server.types import Payload +from cybergym.server.types import Payload, server_conf from cybergym.task.types import verify_task from cybergym.utils import get_arvo_id, get_oss_fuzz_id FLAG = "flag{3xpl017_cyb3r6ym}" DEFAULT_DOCKER_TIMEOUT = 30 # seconds for docker container to run DEFAULT_CMD_TIMEOUT = 10 # seconds for command to run +DEFAULT_RUNNER_IMAGE = "cybergym/oss-fuzz-base-runner:latest" class CustomExitCode(IntEnum): @@ -63,7 +66,6 @@ def is_integer(s): return False - def run_container( task_id: str, poc_path: Path, @@ -99,7 +101,87 @@ def run_container( if container: container.remove(force=True) - return exit_code, docker_output + return exit_code, docker_output + + +def run_container2( + task_id: str, + poc_path: Path, + mode: Literal["vul", "fix"], + data_dir: Path, + docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, + cmd_timeout: int = DEFAULT_CMD_TIMEOUT, +): + client = docker.from_env() + subset, subid = task_id.split(":") + cmd: list[str] + volumes: dict[str, dict[str, str]] + runner_image: str = DEFAULT_RUNNER_IMAGE + container = None + + if subset == "arvo": + runner_image_file = data_dir / "arvo" / subid / "runner_image.txt" + if runner_image_file.exists(): + runner_image = runner_image_file.read_text().strip() + bin_dir = data_dir / "arvo" / subid / mode + volumes = { + str(bin_dir / "arvo"): { + "bind": "/arvo", + "mode": "ro", + }, + str(poc_path.absolute()): { + "bind": "/tmp/poc", # noqa: S108 + "mode": "ro", + }, + str(bin_dir / "libs"): { + "bind": "/out-libs", + "mode": "ro", + }, + } + cmd = ["env", "LD_LIBRARY_PATH=/out-libs", "/bin/bash", "/arvo"] + elif subset == "oss-fuzz": + 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" + with open(meta_file) as f: + metadata = json.load(f) + fuzzer_name = metadata["fuzz_target"] + volumes = {str(poc_path.absolute()): {"bind": "/testcase", "mode": "ro"}} + for subfile in out_dir.iterdir(): + host_path = str(subfile.absolute()) + container_path = os.path.join("/out", subfile.name) + volumes[host_path] = {"bind": container_path, "mode": "ro"} + cmd = ["reproduce", fuzzer_name] + else: + raise HTTPException(status_code=400, detail="Invalid task_id format") + + try: + container = client.containers.run( + image=runner_image, + command=["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} {shlex.join(cmd)} 2>&1"], + detach=True, + volumes=volumes, + ) + out = container.logs(stdout=True, stderr=False, stream=True, follow=True) + exit_code = container.wait(timeout=docker_timeout)["StatusCode"] + if exit_code == 137: # Process killed by timeout + exit_code = CustomExitCode.Timeout + docker_output = b"" + else: + docker_output = b"".join(out) + except requests.exceptions.ReadTimeout: + raise HTTPException(status_code=500, detail="Timeout waiting for the program") from None + except DockerException as e: + raise HTTPException(status_code=500, detail=f"Running error: {e}") from None + except Exception as e: + raise HTTPException(status_code=500, detail=f"Unexpected error: {e}") from None + finally: + if container: + container.remove(force=True) + + return exit_code, docker_output def get_poc_storage_path(poc_id: str, log_dir: Path): @@ -107,7 +189,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): +def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: str, use2: 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") @@ -162,7 +244,10 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st ) # Run the PoC - exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode) + if use2: + exit_code, docker_output = run_container2(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}" with open(output_file, "wb") as f: f.write(docker_output) diff --git a/src/cybergym/server/types.py b/src/cybergym/server/types.py index 54b88ce..41784f8 100644 --- a/src/cybergym/server/types.py +++ b/src/cybergym/server/types.py @@ -1,4 +1,26 @@ -from pydantic import BaseModel +from pathlib import Path + +from pydantic import BaseModel, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from cybergym.task.types import DEFAULT_SALT + + +class ServerConfig(BaseSettings): + salt: str = Field(default=DEFAULT_SALT, description="Salt for checksum") + log_dir: Path = Field(default=Path("./logs"), description="Directory to store logs") + db_path: Path = Field(default=Path("./poc.db"), description="Path to SQLite DB") + binary_dir: Path | None = Field(default=None, description="Directory to store target binaries") + api_key: str = Field( + default="cybergym-030a0cd7-5908-4862-8ab9-91f2bfc7b56d", description="API key for authentication" + ) + api_key_name: str = Field(default="X-API-Key", description="Name of the API key header") + max_file_size_mb: int = Field(default=10, description="Maximum file size for uploads in MB") + + model_config = SettingsConfigDict(env_prefix="CYBERGYM_") + + +server_conf = ServerConfig() class Payload(BaseModel): @@ -15,4 +37,4 @@ class PocQuery(BaseModel): class VerifyPocs(BaseModel): - agent_id: str \ No newline at end of file + agent_id: str From 422c53917cff0532b128a68708d1145de62b3b41 Mon Sep 17 00:00:00 2001 From: zhun Date: Sun, 8 Feb 2026 11:30:43 -0800 Subject: [PATCH 3/5] add base images --- scripts/server_data/binary_only_runners.txt | 6 ++++++ src/cybergym/server/server_utils.py | 17 ++++++++++++++--- 2 files changed, 20 insertions(+), 3 deletions(-) create mode 100644 scripts/server_data/binary_only_runners.txt diff --git a/scripts/server_data/binary_only_runners.txt b/scripts/server_data/binary_only_runners.txt new file mode 100644 index 0000000..3f87062 --- /dev/null +++ b/scripts/server_data/binary_only_runners.txt @@ -0,0 +1,6 @@ +cybergym/oss-fuzz-base-runner:latest +cybergym/oss-fuzz-base-runner:20200102 +cybergym/oss-fuzz-base-runner:20190802 +cybergym/oss-fuzz-base-runner:20200102 +cybergym/oss-fuzz-base-runner:20190802 +cybergym/oss-fuzz-base-runner:20220102 diff --git a/src/cybergym/server/server_utils.py b/src/cybergym/server/server_utils.py index 7e9288f..289e301 100644 --- a/src/cybergym/server/server_utils.py +++ b/src/cybergym/server/server_utils.py @@ -138,6 +138,11 @@ def run_container2( "mode": "ro", }, } + for file in (bin_dir / "out").iterdir(): + volumes[str(file)] = { + "bind": f"/out/{file.name}", + "mode": "ro", + } cmd = ["env", "LD_LIBRARY_PATH=/out-libs", "/bin/bash", "/arvo"] elif subset == "oss-fuzz": if not is_integer(subid): @@ -263,7 +268,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): +def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False, use2: 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") @@ -276,7 +281,10 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False): if rerun or record.vul_exit_code is None: # Run the PoC - exit_code, docker_output = run_container(record.task_id, poc_path, "vul") + if use2: + exit_code, docker_output = run_container2(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: f.write(docker_output) update_poc_output(db, record, "vul", exit_code) @@ -287,7 +295,10 @@ def run_poc_id(db: Session, log_dir: Path, poc_id: str, rerun: bool = False): if rerun or record.fix_exit_code is None: # Run the PoC - exit_code, docker_output = run_container(record.task_id, poc_path, "fix") + if use2: + exit_code, docker_output = run_container2(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: f.write(docker_output) update_poc_output(db, record, "fix", exit_code) From f86dfd7ca94f05ed95a7431a478daf1520d97667 Mon Sep 17 00:00:00 2001 From: NSigma Date: Sun, 8 Feb 2026 15:55:19 -0800 Subject: [PATCH 4/5] Remove duplicate runner versions from binary_only_runners.txt Removed duplicate entries for specific runner versions. --- scripts/server_data/binary_only_runners.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/scripts/server_data/binary_only_runners.txt b/scripts/server_data/binary_only_runners.txt index 3f87062..b43c37f 100644 --- a/scripts/server_data/binary_only_runners.txt +++ b/scripts/server_data/binary_only_runners.txt @@ -1,6 +1,4 @@ cybergym/oss-fuzz-base-runner:latest cybergym/oss-fuzz-base-runner:20200102 cybergym/oss-fuzz-base-runner:20190802 -cybergym/oss-fuzz-base-runner:20200102 -cybergym/oss-fuzz-base-runner:20190802 cybergym/oss-fuzz-base-runner:20220102 From bcbe562ee6247a8bdc64e4cba11167b6fcd66603 Mon Sep 17 00:00:00 2001 From: stneng Date: Mon, 9 Feb 2026 16:18:46 -0800 Subject: [PATCH 5/5] 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: