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