server/environment.py

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