CyberArena / app /sandbox /workdir.py
Hussien Haider
H
80a4a65
Raw
History Blame Contribute Delete
2.61 kB
"""Per-challenge temp workdir for the OS simulator.
Each (team_role, challenge_id) tuple gets its own tempdir; the row's
``files`` JSONB is materialized into real files on first access. The
workdir is cached in-process so the user can `cd` / write scripts
across multiple terminal calls.
"""
import os
import tempfile
import base64
from typing import Optional, Tuple
from app.services.supabase_service import fetch_scenario_by_id
# In-process cache: key -> (workdir_path, last_loaded_row_dict)
_workdirs: dict[str, str] = {}
_last_loaded_row: dict[str, dict] = {}
async def get_or_create_workdir(team_role: str, challenge_id: str) -> Tuple[str, dict]:
"""Return (workdir_path, challenge_row) — cached per challenge.
If the row isn't in Supabase (e.g. ad-hoc training mode), an empty
workdir is created and an empty row dict is returned.
"""
key = f"{team_role}:{challenge_id}"
if key in _workdirs and os.path.isdir(_workdirs[key]):
return _workdirs[key], _last_loaded_row.get(key, {})
row: Optional[dict] = None
if challenge_id:
try:
row = await fetch_scenario_by_id(team_role, challenge_id)
except Exception:
row = None
if row is None:
workdir = tempfile.mkdtemp(prefix=f"ca_{team_role}_{challenge_id[:8]}_")
_workdirs[key] = workdir
return workdir, {}
workdir = tempfile.mkdtemp(prefix=f"ca_{team_role}_{challenge_id[:8]}_")
files = row.get("files") or {}
for filename, b64 in files.items():
# Sanitize: /etc/shadow -> etc/shadow inside workdir
clean = filename.lstrip("/\\").replace("..", "_").replace("\\", "/")
target = os.path.join(workdir, clean)
os.makedirs(os.path.dirname(target) or workdir, exist_ok=True)
try:
data = base64.b64decode(b64)
with open(target, "wb") as f:
f.write(data)
except Exception:
pass
with open(os.path.join(workdir, ".challenge_id"), "w") as f:
f.write(challenge_id)
_workdirs[key] = workdir
_last_loaded_row[key] = row
return workdir, row
def safe_join(workdir: str, filename: str) -> str:
"""Resolve ``filename`` inside ``workdir``; block path traversal."""
clean = (filename or "").lstrip("/\\").replace("..", "_").replace("\\", "/")
if not clean:
raise ValueError("empty filename")
target = os.path.normpath(os.path.join(workdir, clean))
workdir_abs = os.path.normpath(workdir)
if not target.startswith(workdir_abs):
raise ValueError("path traversal blocked")
return target