from __future__ import annotations import json import os import shutil import subprocess import threading import uuid from dataclasses import asdict, dataclass, field from datetime import datetime, timezone from enum import Enum from pathlib import Path from typing import Any from graders.gitleaks_eval import SecurityGrader, git_integrity_message from graders.health_eval import HealthGrader, health_failure_message from graders.reward import STEP_BUDGETS, RewardBreakdown from server.observation import compute_ranked_actions, get_top_blocker def utc_now() -> str: return datetime.now(timezone.utc).isoformat() # --------------------------------------------------------------------------- # Secret visibility tiers # --------------------------------------------------------------------------- class SecretVisibility(str, Enum): SURFACE = "surface" # visible in raw file content on first obs SHALLOW = "shallow" # visible after agent calls inspect_file(path) DEEP = "deep" # visible only after inspect_git_history or inspect_encoded CASCADING = "cascading" # revealed only AFTER another secret is fixed/removed # --------------------------------------------------------------------------- # Dataclasses # --------------------------------------------------------------------------- @dataclass class CommandResult: action: str exit_code: int stdout: str stderr: str timed_out: bool executed_at: str = field(default_factory=utc_now) @dataclass class SessionState: session_id: str task_id: int category: str title: str description: str workspace: str initial_leaks: int current_leaks: int security_score: float health_score: float efficiency_bonus: float reward: float step_budget: int steps_taken: int steps_remaining: int visible_secrets: list[dict[str, Any]] = field(default_factory=list) hidden_count_hint: int = 0 ranked_actions: list[dict[str, Any]] = field(default_factory=list) top_blocker: str = "" conflict_map: dict[str, Any] = field(default_factory=dict) health_stdout: str = "" health_stderr: str = "" health_exit_code: int = 0 observation: str = "" error: str = "none" step_count: int = 0 last_result: CommandResult | None = None created_at: str = field(default_factory=utc_now) updated_at: str = field(default_factory=utc_now) # Internal tracking (not serialised in to_dict output below) _inspected_files: set[str] = field(default_factory=set, repr=False) _deep_inspected: set[str] = field(default_factory=set, repr=False) _fixed_secrets: set[str] = field(default_factory=set, repr=False) def to_dict(self) -> dict[str, Any]: payload = asdict(self) # Remove private tracking fields from API output for private in ("_inspected_files", "_deep_inspected", "_fixed_secrets"): payload.pop(private, None) if self.last_result is None: payload["last_result"] = None return payload # --------------------------------------------------------------------------- # Environment # --------------------------------------------------------------------------- class SecretsAuditEnvironment: def __init__( self, repo_root: Path | None = None, runtime_root: Path | None = None, command_timeout_seconds: int = 90, ) -> None: self.repo_root = (repo_root or Path(__file__).resolve().parents[1]).resolve() self.tasks_root = self.repo_root / "tasks" self.runtime_root = (runtime_root or self.repo_root / "runtime").resolve() self.runtime_root.mkdir(parents=True, exist_ok=True) self.command_timeout_seconds = command_timeout_seconds self.security_grader = SecurityGrader() self.health_grader = HealthGrader() self._lock = threading.RLock() self._session: SessionState | None = None self._task_meta: dict[str, Any] | None = None self._workspace_path: Path | None = None # ------------------------------------------------------------------ # # Public API # # ------------------------------------------------------------------ # def list_tasks(self) -> list[dict[str, Any]]: tasks: list[dict[str, Any]] = [] for task_file in sorted(self.tasks_root.glob("*/task_*/task.json")): with task_file.open("r", encoding="utf-8") as handle: task_data = json.load(handle) # Add task_id field for OpenEnv compatibility task_num = task_data.get("id", 0) task_data["task_id"] = f"task_{task_num}" tasks.append(task_data) return tasks def state(self) -> dict[str, Any]: with self._lock: payload = { "session": None if self._session is None else self._session.to_dict(), "available_tasks": self.list_tasks(), } return self._mask_value(payload) def reset(self, task_id: int) -> dict[str, Any]: with self._lock: task_dir = self._locate_task(task_id) meta = self._load_task_meta(task_dir) session_id = str(uuid.uuid4()) workspace = self.runtime_root / f"session_{session_id}" workspace.parent.mkdir(parents=True, exist_ok=True) self._materialize_task(task_dir, meta, workspace) if not (workspace / ".git").exists(): self._initialize_git_workspace(workspace) workspace = workspace.resolve() initial_security = self.security_grader.grade(workspace, meta) initial_health = self.health_grader.grade(workspace) difficulty = meta.get("category", "medium") step_budget = STEP_BUDGETS.get(difficulty, STEP_BUDGETS["medium"]) reward = RewardBreakdown.from_reports( initial_security, initial_health, steps_taken=0, difficulty=difficulty, ) # Build initial visible secrets all_secrets: list[dict[str, Any]] = meta.get("secrets", []) known_files = list({s["path"] for s in all_secrets}) empty_set: set[str] = set() visible = self._compute_visible_secrets(all_secrets, empty_set, empty_set, empty_set) hidden_count = len(all_secrets) - len(visible) obs_text = self._build_observation(initial_health, initial_security, None) ranked = compute_ranked_actions( inspected_files=empty_set, deep_inspected=empty_set, fixed_secrets=empty_set, visible_secrets=visible, known_files=known_files, difficulty=difficulty, hidden_count=hidden_count, ) blocker = get_top_blocker(visible, empty_set, hidden_count, difficulty) conflict_map = self._build_conflict_map(meta, visible) self._task_meta = meta self._workspace_path = workspace self._session = SessionState( session_id=session_id, task_id=meta["id"], category=difficulty, title=meta["title"], description=meta.get("description", ""), workspace=".", initial_leaks=initial_security.finding_count, current_leaks=initial_security.finding_count, security_score=reward.security_score, health_score=reward.health_score, efficiency_bonus=reward.efficiency_bonus, reward=reward.total_reward, step_budget=step_budget, steps_taken=0, steps_remaining=step_budget, visible_secrets=visible, hidden_count_hint=hidden_count, ranked_actions=ranked, top_blocker=blocker, conflict_map=conflict_map, health_stdout=initial_health.stdout, health_stderr=initial_health.stderr, health_exit_code=initial_health.exit_code, observation=obs_text, error=self._build_error(initial_health, initial_security, None), _inspected_files=set(), _deep_inspected=set(), _fixed_secrets=set(), ) self._session.updated_at = utc_now() return self.state() def step(self, action: str) -> dict[str, Any]: with self._lock: if self._session is None or self._task_meta is None or self._workspace_path is None: raise RuntimeError("No active session. Call /reset first.") workspace = self._workspace_path session = self._session meta = self._task_meta difficulty = meta.get("category", "medium") step_budget = STEP_BUDGETS.get(difficulty, STEP_BUDGETS["medium"]) # ---- Parse structured actions first ------------------------- structured_result = self._handle_structured_action(action, session) if structured_result is not None: exit_code, stdout, stderr, timed_out = structured_result else: # Fall through to raw bash execution timed_out = False try: completed = subprocess.run( ["/usr/bin/bash", "-lc", action], cwd=workspace, capture_output=True, text=True, timeout=self.command_timeout_seconds, ) exit_code = completed.returncode stdout = completed.stdout stderr = completed.stderr except subprocess.TimeoutExpired as exc: timed_out = True exit_code = 124 stdout = exc.stdout or "" stderr = exc.stderr or f"Command timed out after {self.command_timeout_seconds}s." session.step_count += 1 session.last_result = CommandResult( action=action, exit_code=exit_code, stdout=stdout, stderr=stderr, timed_out=timed_out, ) # ---- Re-grade ----------------------------------------------- security = self.security_grader.grade(workspace, meta) health = self.health_grader.grade(workspace) reward = RewardBreakdown.from_reports( security, health, initial_leaks=session.initial_leaks, steps_taken=session.step_count, difficulty=difficulty, ) # ---- Infer fixed secrets (leaks dropped → mark surface secrets fixed) ---- self._update_fixed_secrets(session, meta, security) # ---- Recompute visible secrets & obs fields ------------------ all_secrets = meta.get("secrets", []) known_files = list({s["path"] for s in all_secrets}) visible = self._compute_visible_secrets( all_secrets, session._inspected_files, session._deep_inspected, session._fixed_secrets, ) hidden_count = len(all_secrets) - len(visible) ranked = compute_ranked_actions( inspected_files=session._inspected_files, deep_inspected=session._deep_inspected, fixed_secrets=session._fixed_secrets, visible_secrets=visible, known_files=known_files, difficulty=difficulty, hidden_count=hidden_count, ) blocker = get_top_blocker(visible, session._fixed_secrets, hidden_count, difficulty) conflict_map = self._build_conflict_map(meta, visible) steps_taken = session.step_count session.current_leaks = security.finding_count session.security_score = reward.security_score session.health_score = reward.health_score session.efficiency_bonus = reward.efficiency_bonus session.reward = reward.total_reward session.steps_taken = steps_taken session.steps_remaining = max(0, step_budget - steps_taken) session.visible_secrets = visible session.hidden_count_hint = hidden_count session.ranked_actions = ranked session.top_blocker = blocker session.conflict_map = conflict_map session.health_stdout = health.stdout session.health_stderr = health.stderr session.health_exit_code = health.exit_code session.observation = self._build_observation(health, security, session.last_result) session.error = self._build_error(health, security, session.last_result) session.updated_at = utc_now() return self.state() # ------------------------------------------------------------------ # # Structured action handlers # # ------------------------------------------------------------------ # def _handle_structured_action( self, action: str, session: SessionState, ) -> tuple[int, str, str, bool] | None: """Parse special prefixed actions. Returns (exit_code, stdout, stderr, timed_out) or None if the action should fall through to bash execution.""" stripped = action.strip() if stripped.startswith("inspect_file "): path = stripped[len("inspect_file "):].strip() session._inspected_files.add(path) return 0, f"[inspect_file] Marked '{path}' as inspected. SHALLOW secrets now visible.", "", False if stripped.startswith("inspect_git_history"): parts = stripped.split(maxsplit=1) path = parts[1].strip() if len(parts) > 1 else "." session._deep_inspected.add(path) return 0, f"[inspect_git_history] Scanned git history at '{path}'. DEEP secrets now visible.", "", False if stripped.startswith("inspect_encoded "): rest = stripped[len("inspect_encoded "):].strip().split() path = rest[0] if rest else "." line = rest[1] if len(rest) > 1 else "0" key = f"{path}:{line}" session._deep_inspected.add(key) session._deep_inspected.add(path) # also unlock by path alone return 0, f"[inspect_encoded] Decoded blob at '{path}:{line}'. DEEP encoded secrets now visible.", "", False return None # ------------------------------------------------------------------ # # Visibility computation # # ------------------------------------------------------------------ # def _compute_visible_secrets( self, all_secrets: list[dict[str, Any]], inspected_files: set[str], deep_inspected: set[str], fixed_secrets: set[str], ) -> list[dict[str, Any]]: """Return the subset of secrets the agent can currently see.""" visible: list[dict[str, Any]] = [] for s in all_secrets: vis = s.get("visibility", SecretVisibility.SURFACE) path = s.get("path", "") trigger = s.get("trigger_secret_id") if vis == SecretVisibility.SURFACE: visible.append(s) elif vis == SecretVisibility.SHALLOW: if path in inspected_files: visible.append(s) elif vis == SecretVisibility.DEEP: req = s.get("requires_action", "") if req == "inspect_encoded": # Check path or path:line in deep_inspected if path in deep_inspected or any(k.startswith(path + ":") for k in deep_inspected): visible.append(s) else: # inspect_git_history: any deep_inspected entry unlocks it if deep_inspected: visible.append(s) elif vis == SecretVisibility.CASCADING: if trigger and trigger in fixed_secrets: visible.append(s) return visible def _update_fixed_secrets( self, session: SessionState, meta: dict[str, Any], security_report: Any, ) -> None: """Heuristically infer which secrets are fixed based on leak count drop. Mark all currently visible surface secrets as fixed if the total finding_count dropped to 0. For more granular tracking we compare current_leaks vs initial_leaks and mark the first unfixed visible surface secret. """ all_secrets = meta.get("secrets", []) if security_report.finding_count == 0: for s in all_secrets: session._fixed_secrets.add(s["id"]) elif security_report.finding_count < session.current_leaks: # Some leaks were fixed — mark unresolved surface secrets iteratively visible_unfixed = [ s for s in self._compute_visible_secrets( all_secrets, session._inspected_files, session._deep_inspected, session._fixed_secrets, ) if s["id"] not in session._fixed_secrets and s.get("visibility") == SecretVisibility.SURFACE ] for s in visible_unfixed: session._fixed_secrets.add(s["id"]) break # mark one per step # ------------------------------------------------------------------ # # Conflict map # # ------------------------------------------------------------------ # def _build_conflict_map( self, meta: dict[str, Any], visible_secrets: list[dict[str, Any]], ) -> dict[str, Any]: """Return conflict map filtered to currently-visible secrets only.""" raw_map = meta.get("conflict_map", {}) if not raw_map: return {} visible_ids = {s["id"] for s in visible_secrets} result: dict[str, Any] = {} for secret_id, conflicts in raw_map.items(): if secret_id in visible_ids: result[secret_id] = { "reveals": conflicts.get("reveals", []), "blocks": conflicts.get("blocks", []), "note": conflicts.get("note", ""), } return result # ------------------------------------------------------------------ # # Task management # # ------------------------------------------------------------------ # def _locate_task(self, task_id: int) -> Path: matches = list(self.tasks_root.glob(f"*/task_{task_id:02d}")) if not matches: raise FileNotFoundError(f"Task {task_id} was not found under {self.tasks_root}.") return matches[0] def _load_task_meta(self, task_dir: Path) -> dict[str, Any]: with (task_dir / "task.json").open("r", encoding="utf-8") as handle: return json.load(handle) def _materialize_task(self, task_dir: Path, meta: dict[str, Any], workspace: Path) -> None: if workspace.exists(): shutil.rmtree(workspace) build_script = meta.get("build_script") if build_script: workspace.mkdir(parents=True, exist_ok=True) subprocess.run( [ "python3", str(task_dir / build_script), "--output", str(workspace), ], cwd=task_dir, check=True, ) return shutil.copytree(task_dir, workspace, ignore=shutil.ignore_patterns("task.json", "build_repo.py")) def _initialize_git_workspace(self, workspace: Path) -> None: commands = [ ["git", "init"], ["git", "config", "user.name", "SecretsAuditEnv"], ["git", "config", "user.email", "benchmark@example.com"], ["git", "add", "."], ["git", "commit", "-m", "Initial vulnerable state"], ] for command in commands: subprocess.run(command, cwd=workspace, check=True, capture_output=True, text=True) # ------------------------------------------------------------------ # # Observation building # # ------------------------------------------------------------------ # def _build_observation(self, health: Any, security: Any, result: CommandResult | None) -> str: parts: list[str] = [] if result is not None: if result.stdout.strip(): parts.append(f"Command stdout:\n{result.stdout.strip()}") if result.stderr.strip(): parts.append(f"Command stderr:\n{result.stderr.strip()}") health_message = health_failure_message(health) if health_message: parts.append(f"Health check output:\n{health_message}") git_message = git_integrity_message(security) if git_message: parts.append(git_message) if not parts: return "none" return "\n\n".join(parts) def _build_error(self, health: Any, security: Any, result: CommandResult | None) -> str: if result is not None: if result.timed_out: return result.stderr.strip() or "Command timed out." if result.exit_code != 0: return result.stderr.strip() or result.stdout.strip() or "Command failed." health_message = health_failure_message(health) if health_message: return health_message git_message = git_integrity_message(security) if git_message: return git_message return "none" # ------------------------------------------------------------------ # # Path masking # # ------------------------------------------------------------------ # def _mask_value(self, value: Any) -> Any: if isinstance(value, dict): return {key: self._mask_value(item) for key, item in value.items()} if isinstance(value, list): return [self._mask_value(item) for item in value] if isinstance(value, str): return self._mask_text(value) return value def _mask_text(self, text: str) -> str: if not text: return text masked = text candidates = {str(self.repo_root), os.path.expanduser("~")} if self._workspace_path is not None: candidates.add(str(self._workspace_path)) for candidate in sorted((item for item in candidates if item), key=len, reverse=True): masked = masked.replace(f"{candidate}/", "") masked = masked.replace(candidate, ".") return masked