File size: 2,614 Bytes
80a4a65
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
"""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