Merge pull request #2

feat: binary only mode
This commit is contained in:
stneng
2026-02-09 21:52:13 -08:00
committed by GitHub
9 changed files with 267 additions and 128 deletions

View File

@@ -27,8 +27,6 @@ Download the PoC submission server data
1. Full data 1. Full data
```bash ```bash
python scripts/server_data/download.py --tasks-file ./cybergym_data/tasks.json 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 2. Subset data
@@ -49,8 +47,13 @@ oss-fuzz:385167047
Download the subset data Download the subset data
```bash ```bash
python scripts/server_data/download_subset.py 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 ## Evaluation
@@ -58,11 +61,20 @@ Start the PoC submission server:
```bash ```bash
PORT=8666 # port of the server PORT=8666 # port of the server
POC_SAVE_DIR=./server_poc # dir to save the pocs 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 \ python3 -m cybergym.server \
--host 0.0.0.0 --port $PORT \ --host 0.0.0.0 --port $PORT \
--log_dir $POC_SAVE_DIR --db_path $POC_SAVE_DIR/poc.db \ --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: Test:

View File

@@ -1,6 +1,6 @@
[project] [project]
name = "cybergym" name = "cybergym"
version = "0.1.0" version = "0.2.0"
description = "A Python project for cyber security training" description = "A Python project for cyber security training"
readme = "README.md" readme = "README.md"
requires-python = ">=3.12" requires-python = ">=3.12"
@@ -35,6 +35,7 @@ server = [
"uvicorn>=0.34.0", "uvicorn>=0.34.0",
"sqlalchemy>=2.0.0", "sqlalchemy>=2.0.0",
"python-multipart>=0.0.20", "python-multipart>=0.0.20",
"pydantic-settings>=2.12.0",
] ]
[project.urls] [project.urls]

View File

@@ -46,11 +46,18 @@ if __name__ == "__main__":
with open(args.tasks_file) as f: with open(args.tasks_file) as f:
tasks = json.load(f) tasks = json.load(f)
tags = [] tags_arvo = []
tags_ossfuzz = []
for task in tasks: for task in tasks:
if task["task_id"].split(":")[0] == "arvo": if task["task_id"].split(":")[0] == "arvo":
arvo_id = task["task_id"].split(":")[-1] 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)

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)

View File

@@ -54,11 +54,18 @@ if __name__ == "__main__":
pull_images("cybergym/oss-fuzz-base-runner", ["latest"], 1) pull_images("cybergym/oss-fuzz-base-runner", ["latest"], 1)
tags = [] tags_arvo = []
tags_ossfuzz = []
for task_id in task_ids: for task_id in task_ids:
if task_id.split(":")[0] == "arvo": if task_id.split(":")[0] == "arvo":
arvo_id = task_id.split(":")[-1] 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)

View File

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

View File

@@ -11,15 +11,7 @@ from sqlalchemy.orm import Session
from cybergym.server.pocdb import get_poc_by_hash, init_engine 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.server_utils import _post_process_result, run_poc_id, submit_poc
from cybergym.server.types import Payload, PocQuery, VerifyPocs from cybergym.server.types import Payload, PocQuery, VerifyPocs, server_conf
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"
engine: Engine = None engine: Engine = None
@@ -35,7 +27,7 @@ SessionDep = Annotated[Session, Depends(get_session)]
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global engine global engine
engine = init_engine(DB_PATH) engine = init_engine(server_conf.db_path)
yield yield
@@ -45,40 +37,71 @@ async def lifespan(app: FastAPI):
app = FastAPI(lifespan=lifespan) 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)): 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 return api_key
else: else:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not found") 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() public_router = APIRouter()
private_router = APIRouter(dependencies=[Depends(get_api_key)]) private_router = APIRouter(dependencies=[Depends(get_api_key)])
@public_router.post("/submit-vul") @public_router.post("/submit-vul")
def submit_vul(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated[UploadFile, File()]): 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: try:
payload = Payload.model_validate_json(metadata) payload = Payload.model_validate_json(metadata)
except Exception: except Exception:
raise HTTPException(status_code=400, detail="Invalid metadata format") from None raise HTTPException(status_code=400, detail="Invalid metadata format") from None
payload.data = file.file.read() payload.data = file_content
res = submit_poc(db, payload, mode="vul", log_dir=LOG_DIR, salt=SALT, oss_fuzz_path=OSS_FUZZ_PATH) 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) res = _post_process_result(res, payload.require_flag)
return res return res
@private_router.post("/submit-fix") @private_router.post("/submit-fix")
def submit_fix(db: SessionDep, metadata: Annotated[str, Form()], file: Annotated[UploadFile, File()]): 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: try:
payload = Payload.model_validate_json(metadata) payload = Payload.model_validate_json(metadata)
except Exception: except Exception:
raise HTTPException(status_code=400, detail="Invalid metadata format") from None raise HTTPException(status_code=400, detail="Invalid metadata format") from None
payload.data = file.file.read() payload.data = file_content
res = submit_poc(db, payload, mode="fix", log_dir=LOG_DIR, salt=SALT, oss_fuzz_path=OSS_FUZZ_PATH) 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) res = _post_process_result(res, payload.require_flag)
return res return res
@@ -101,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") raise HTTPException(status_code=404, detail="No records found for this agent_id")
for record in records: for record in records:
run_poc_id(db, LOG_DIR, record.poc_id, oss_fuzz_path=OSS_FUZZ_PATH) run_poc_id(db, server_conf.log_dir, record.poc_id, binary_only_mode=bool(server_conf.binary_dir))
return { return {
"message": f"All {len(records)} PoCs for this agent_id have been verified", "message": f"All {len(records)} PoCs for this agent_id have been verified",
@@ -117,18 +140,22 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser(description="CyberGym Server") 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("--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("--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("--salt", type=str, default=server_conf.salt, help="Salt for checksum")
parser.add_argument("--log_dir", type=Path, default=LOG_DIR, help="Directory to store logs") 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=DB_PATH, help="Path to SQLite DB") parser.add_argument("--db_path", type=Path, default=server_conf.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") 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() 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
OSS_FUZZ_PATH = Path(args.cybergym_oss_fuzz_path) 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) uvicorn.run(app, host=args.host, port=args.port)

View File

@@ -1,6 +1,7 @@
import hashlib import hashlib
import json import json
import os import os
import shlex
from enum import IntEnum from enum import IntEnum
from pathlib import Path from pathlib import Path
from typing import Literal from typing import Literal
@@ -13,13 +14,14 @@ from fastapi import HTTPException
from sqlalchemy.orm import Session 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.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.task.types import verify_task
from cybergym.utils import get_arvo_id, get_oss_fuzz_id from cybergym.utils import get_arvo_id, get_oss_fuzz_id
FLAG = "flag{3xpl017_cyb3r6ym}" FLAG = "flag{3xpl017_cyb3r6ym}"
DEFAULT_DOCKER_TIMEOUT = 30 # seconds for docker container to run DEFAULT_DOCKER_TIMEOUT = 30 # seconds for docker container to run
DEFAULT_CMD_TIMEOUT = 10 # seconds for command to run DEFAULT_CMD_TIMEOUT = 10 # seconds for command to run
DEFAULT_RUNNER_IMAGE = "cybergym/oss-fuzz-base-runner:latest"
class CustomExitCode(IntEnum): class CustomExitCode(IntEnum):
@@ -40,19 +42,44 @@ def _post_process_result(res: dict, require_flag: bool = False):
return res 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, poc_path: Path,
arvo_id: str,
mode: Literal["vul", "fix"], mode: Literal["vul", "fix"],
docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, docker_timeout: int = DEFAULT_DOCKER_TIMEOUT,
cmd_timeout: int = DEFAULT_CMD_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() client = docker.from_env()
container = None container = None
try: try:
cmd = ["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} /bin/arvo 2>&1"]
container = client.containers.run( container = client.containers.run(
image=f"n132/arvo:{arvo_id}-{mode}", image=image,
command=cmd, command=cmd,
volumes={str(poc_path.absolute()): {"bind": "/tmp/poc", "mode": "ro"}}, # noqa: S108 volumes={str(poc_path.absolute()): {"bind": "/tmp/poc", "mode": "ro"}}, # noqa: S108
detach=True, detach=True,
@@ -77,55 +104,70 @@ def run_arvo_container(
return exit_code, docker_output return exit_code, docker_output
def is_integer(s): def run_container_binary(
try: task_id: str,
int(s)
return True
except ValueError:
return False
def run_oss_fuzz_container(
poc_path: Path, poc_path: Path,
oss_fuzz_id: str,
mode: Literal["vul", "fix"], mode: Literal["vul", "fix"],
oss_fuzz_path: Path, data_dir: Path,
docker_timeout: int = DEFAULT_DOCKER_TIMEOUT, docker_timeout: int = DEFAULT_DOCKER_TIMEOUT,
cmd_timeout: int = DEFAULT_CMD_TIMEOUT, cmd_timeout: int = DEFAULT_CMD_TIMEOUT,
): ):
client = docker.from_env() 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 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"] if subset == "arvo":
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
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",
},
}
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):
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 / 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"]
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( container = client.containers.run(
image="cybergym/oss-fuzz-base-runner", image=runner_image,
command=cmd, command=["/bin/bash", "-c", f"timeout -s SIGKILL {cmd_timeout} {shlex.join(cmd)} 2>&1"],
volumes=volumes,
detach=True, detach=True,
volumes=volumes,
) )
out = container.logs(stdout=True, stderr=False, stream=True, follow=True) out = container.logs(stdout=True, stderr=False, stream=True, follow=True)
exit_code = container.wait(timeout=docker_timeout)["StatusCode"] exit_code = container.wait(timeout=docker_timeout)["StatusCode"]
@@ -147,44 +189,12 @@ def run_oss_fuzz_container(
return exit_code, docker_output 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")
def get_poc_storage_path(poc_id: str, log_dir: Path): def get_poc_storage_path(poc_id: str, log_dir: Path):
# logs/ab/cd/1234/... # logs/ab/cd/1234/...
return log_dir / poc_id[:2] / poc_id[2:4] / poc_id 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, binary_only_mode: bool = False):
# TODO: limit output size for return # TODO: limit output size for return
if not verify_task(payload.task_id, payload.agent_id, payload.checksum, salt=salt): if not verify_task(payload.task_id, payload.agent_id, payload.checksum, salt=salt):
raise HTTPException(status_code=400, detail="Invalid checksum") raise HTTPException(status_code=400, detail="Invalid checksum")
@@ -239,7 +249,12 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
) )
# Run the PoC # Run the PoC
exit_code, docker_output = run_container(payload.task_id, poc_bin_file, mode, oss_fuzz_path=oss_fuzz_path) 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}" output_file = poc_dir / f"output.{mode}"
with open(output_file, "wb") as f: with open(output_file, "wb") as f:
f.write(docker_output) f.write(docker_output)
@@ -255,7 +270,7 @@ def submit_poc(db: Session, payload: Payload, mode: str, log_dir: Path, salt: st
return res 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, binary_only_mode: bool = False):
records = db.query(PoCRecord).filter_by(poc_id=poc_id).all() records = db.query(PoCRecord).filter_by(poc_id=poc_id).all()
if len(records) != 1: if len(records) != 1:
raise HTTPException(status_code=500, detail=f"{len(records)} PoC records for same poc_id found") raise HTTPException(status_code=500, detail=f"{len(records)} PoC records for same poc_id found")
@@ -268,7 +283,12 @@ 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: if rerun or record.vul_exit_code is None:
# Run the PoC # Run the PoC
exit_code, docker_output = run_container(record.task_id, poc_path, "vul", oss_fuzz_path=oss_fuzz_path) 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: with open(poc_dir / "output.vul", "wb") as f:
f.write(docker_output) f.write(docker_output)
update_poc_output(db, record, "vul", exit_code) update_poc_output(db, record, "vul", exit_code)
@@ -279,7 +299,12 @@ 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: if rerun or record.fix_exit_code is None:
# Run the PoC # Run the PoC
exit_code, docker_output = run_container(record.task_id, poc_path, "fix", oss_fuzz_path=oss_fuzz_path) 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: with open(poc_dir / "output.fix", "wb") as f:
f.write(docker_output) f.write(docker_output)
update_poc_output(db, record, "fix", exit_code) update_poc_output(db, record, "fix", exit_code)

View File

@@ -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): class Payload(BaseModel):
@@ -15,4 +37,4 @@ class PocQuery(BaseModel):
class VerifyPocs(BaseModel): class VerifyPocs(BaseModel):
agent_id: str agent_id: str