"""Environment logic for long-horizon, partially observable incident response.""" from __future__ import annotations import random import uuid import os from typing import Any from models import ( AgenticSecurityLabAction, AgenticSecurityLabObservation, AgenticSecurityLabState, ) from scenarios import generate_scenario, get_scenario from training.rewards import benchmark_score_from_breakdown, uncertainty_reduction_score # Attempt to use the procedural generator for richer training-mode variation. try: from training.procedural_scenarios import generate_procedural_scenario as _gen_proc _HAS_PROCEDURAL = True except ImportError: _HAS_PROCEDURAL = False # Inherit from OpenEnv's base Environment class if available; fall back to # plain object only when explicitly enabled for local tests. try: from openenv.core import Environment as _OpenEnvBase # type: ignore[import] _HAS_OPENENV_BASE = True except ImportError: _OpenEnvBase = object # type: ignore[misc,assignment] _HAS_OPENENV_BASE = False VALID_COMMANDS = { "inspect_package", "check_dependents", "rotate_secret", "quarantine", "notify", "scan_logs", "conclude", } COMMAND_ALIASES = { "inspect": "inspect_package", "deps": "check_dependents", "dependents": "check_dependents", "rotate": "rotate_secret", } VALID_MODES = {"benchmark", "training"} class AgenticSecurityLabEnvironment(_OpenEnvBase): """Supply-chain incident response environment. Inherits from ``openenv.core.Environment`` when openenv-core is installed, giving the framework standard lifecycle hooks. Falls back to plain ``object`` so the class works in lightweight / testing contexts too. """ def __init__(self, task_name: str = "easy") -> None: allow_fallback = os.getenv("AGENTIC_SECURITY_LAB_ALLOW_ENV_BASE_FALLBACK", "false").lower() == "true" if not _HAS_OPENENV_BASE and not allow_fallback: raise RuntimeError( "openenv-core is required in production. Install openenv-core or set " "AGENTIC_SECURITY_LAB_ALLOW_ENV_BASE_FALLBACK=true for local fallback tests." ) self._task_name = task_name self._state = AgenticSecurityLabState() self._scenario: dict[str, Any] = {} self._rng = random.Random(0) self._last_uncertainty = 1.0 def reset( self, task_name: str | None = None, mode: str | None = None, command_fallback_enabled: bool | None = None, ) -> AgenticSecurityLabObservation: if task_name: self._task_name = task_name requested_mode = (mode or "benchmark").strip().lower() resolved_mode = requested_mode if requested_mode in VALID_MODES else "benchmark" mode_fallback_used = resolved_mode != requested_mode if command_fallback_enabled is None: command_fallback_enabled = False if resolved_mode == "benchmark": scenario = get_scenario(self._task_name) self._rng = random.Random(self._task_name) else: # Prefer fully procedural generation (infinite unique incidents) when # available; fall back to jitter-only variation for minimal installs. if _HAS_PROCEDURAL: scenario = _gen_proc(difficulty=self._task_name) else: scenario = generate_scenario(self._task_name, difficulty_scale=1.0) self._rng = random.Random() self._scenario = scenario self._state = AgenticSecurityLabState( episode_id=str(uuid.uuid4()), step_count=0, task_name=self._task_name, mode=resolved_mode, mode_fallback_used=mode_fallback_used, command_fallback_enabled=bool(command_fallback_enabled), packages=scenario["packages"], dependents=scenario["dependents"], secrets=scenario["secrets"], max_steps=scenario["max_steps"], exfiltration_step=scenario["exfiltration_step"], pending_hidden_iocs=list(scenario.get("hidden_iocs", [])), plan_progress={ "investigate": False, "trace_root_cause": False, "contain": False, "recover": False, "notify": False, "conclude": False, }, ) self._last_uncertainty = 1.0 return self._build_obs( reward=0.0, success=True, result=( f"[INCIDENT ALERT] Task: {self._task_name.upper()}\n" f"{scenario['description']}\n\n" f"Packages in scope: {list(scenario['packages'].keys())}\n" f"Exfiltration window: {scenario['exfiltration_step']} steps " f"(budget: {scenario['max_steps']} total)\n" "Commands: inspect_package, check_dependents, rotate_secret, " "quarantine, notify, scan_logs, conclude" ), steps_left=scenario["exfiltration_step"], data={ "packages_in_scope": list(scenario["packages"].keys()), "mode": resolved_mode, "command_fallback_enabled": bool(command_fallback_enabled), "max_steps": scenario["max_steps"], }, ) def step(self, action: AgenticSecurityLabAction) -> AgenticSecurityLabObservation: state = self._state if state.incident_contained or state.attacker_succeeded or state.step_count >= state.max_steps: return self._terminal_obs(0.0, result="Episode already ended.") state.step_count += 1 steps_left = max(0, state.exfiltration_step - state.step_count) self._advance_attacker() command, fallback_used = self._canonicalize_command(action.command) params = action.parameters if fallback_used: state.command_fallback_used_count += 1 if command not in VALID_COMMANDS: state.invalid_action_count += 1 reward = -0.01 state.total_reward += reward return self._build_obs( reward=reward, success=False, result=f"Unknown command '{action.command}'. Valid: {sorted(VALID_COMMANDS)}", steps_left=steps_left, error=f"Invalid command: {action.command}", ) dispatch = { "inspect_package": self._cmd_inspect, "check_dependents": self._cmd_check_dependents, "rotate_secret": self._cmd_rotate_secret, "quarantine": self._cmd_quarantine, "notify": self._cmd_notify, "scan_logs": self._cmd_scan_logs, "conclude": self._cmd_conclude, } observation = dispatch[command](params, steps_left) self._log_transition(command, params, observation) return observation @property def state(self) -> AgenticSecurityLabState: return self._state def _canonicalize_command(self, raw_command: str) -> tuple[str, bool]: command = raw_command.strip().lower() if command in VALID_COMMANDS: return command, False if self._state.command_fallback_enabled and command in COMMAND_ALIASES: return COMMAND_ALIASES[command], True return command, False def _cmd_inspect(self, params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: package = params.get("package", "") if package not in self._state.packages: return self._invalid_target("Package", package, steps_left) metadata = self._state.packages[package] self._state.inspected.append(package) self._state.plan_progress["investigate"] = True if metadata.get("malicious"): self._remember_package(package) discovered_deps = metadata.get("deps", []) if discovered_deps: self._state.plan_progress["trace_root_cause"] = True inspect_count = self._state.inspected.count(package) reward = max(0.0025, round(0.01 * (0.55 ** max(0, inspect_count - 1)), 4)) self._state.total_reward += reward return self._build_obs( reward=reward, success=True, result=( f"Inspection: {package}\n" f"Publisher: {metadata.get('publisher', 'unknown')}\n" f"Published: {metadata.get('publish_date', 'unknown')}\n" f"Versions: {metadata.get('versions', [])}\n" f"Dependencies: {discovered_deps or 'None'}\n" f"IOC summary: {metadata.get('iocs') or 'No direct IOC found.'}" ), steps_left=steps_left, data={"package": package, "metadata": metadata}, ) def _cmd_check_dependents(self, params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: package = params.get("package", "") dependents = self._state.dependents.get(package) if dependents is None: return self._invalid_target("Package", package, steps_left) self._state.traced_packages.append(package) self._state.plan_progress["trace_root_cause"] = True if self._state.packages.get(package, {}).get("malicious"): self._remember_package(package) reward = 0.01 self._state.total_reward += reward return self._build_obs( reward=reward, success=True, result=f"Dependents of {package} ({len(dependents)} total): {dependents}", steps_left=steps_left, data={"package": package, "dependents": dependents}, ) def _cmd_rotate_secret(self, params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: secret = params.get("secret", "") if secret not in self._state.secrets: return self._invalid_target("Secret", secret, steps_left) secret_meta = self._state.secrets[secret] if secret_meta["rotated"]: reward = -0.02 self._state.total_reward += reward return self._build_obs( reward=reward, success=False, result=f"Secret '{secret}' was already rotated.", steps_left=steps_left, ) secret_meta["rotated"] = True self._state.rotated_secrets.append(secret) self._state.plan_progress["recover"] = True reward = 0.12 if secret_meta["critical"] else 0.06 self._state.total_reward += reward return self._build_obs( reward=reward, success=True, result=( f"Rotated '{secret}' " f"({'CRITICAL' if secret_meta['critical'] else 'standard'}).\n" f"Owner: {secret_meta['owner']}. Old value invalidated." ), steps_left=steps_left, data={"secret": secret, "critical": secret_meta["critical"]}, ) def _cmd_quarantine(self, params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: package = params.get("package", "") if package not in self._state.packages: return self._invalid_target("Package", package, steps_left) if package in self._state.quarantined: return self._build_obs( reward=-0.02, success=False, result=f"'{package}' already quarantined.", steps_left=steps_left, ) package_meta = self._state.packages[package] if not package_meta["malicious"]: reward = -0.05 self._state.false_positive_count += 1 self._state.total_reward += reward return self._build_obs( reward=reward, success=False, result=f"False positive: '{package}' is not malicious.", steps_left=steps_left, ) self._state.quarantined.append(package) self._state.plan_progress["contain"] = True self._remember_package(package) reward = 0.15 self._state.total_reward += reward return self._build_obs( reward=reward, success=True, result=f"Quarantined '{package}'. Registry blocks future installs.", steps_left=steps_left, data={"package": package}, ) def _cmd_notify(self, params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: team = params.get("team", "") valid_teams = set() for teams in self._state.dependents.values(): valid_teams.update(teams) if team not in valid_teams: return self._invalid_target("Team", team, steps_left) if team in self._state.notified_teams: return self._build_obs( reward=0.0, success=True, result=f"Team '{team}' already notified.", steps_left=steps_left, ) self._state.notified_teams.append(team) self._state.plan_progress["notify"] = True reward = 0.04 self._state.total_reward += reward return self._build_obs( reward=reward, success=True, result=f"Notified '{team}' with incident guidance.", steps_left=steps_left, data={"team": team}, ) def _cmd_scan_logs(self, params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: package = params.get("package", "") if package not in self._state.packages: return self._invalid_target("Package", package, steps_left) self._state.scanned_logs.append(package) hints = self._scenario.get("scan_logs_hints", {}) log_result = hints.get(package, f"No suspicious entries found for {package}.") if self._state.packages[package]["malicious"]: self._remember_package(package) self._discover_secrets_for_package(package) if ( self._state.mode == "training" and self._state.pending_hidden_iocs and self._rng.random() < self._scenario.get("stochastic", {}).get("alert_reveal_chance", 0.35) ): hidden = self._state.pending_hidden_iocs.pop(0) self._state.discovered_iocs.append(hidden) log_result = f"{log_result}\nAdditional hidden signal: {hidden}" self._state.plan_progress["investigate"] = True scan_count = self._state.scanned_logs.count(package) reward = max(0.0015, round(0.02 * (0.5 ** max(0, scan_count - 1)), 4)) self._state.total_reward += reward return self._build_obs( reward=reward, success=True, result=f"CI/CD log scan - {package}:\n {log_result}", steps_left=steps_left, data={"package": package, "log_excerpt": log_result}, ) def _cmd_conclude(self, _params: dict[str, Any], steps_left: int) -> AgenticSecurityLabObservation: ratios = self._score_breakdown() q_done = ratios["quarantine_ratio"] == 1.0 r_done = ratios["rotate_ratio"] == 1.0 n_done = ratios["notify_ratio"] == 1.0 contained = ratios["contain_ratio"] == 1.0 and q_done and r_done bonus = 0.0 if q_done: bonus += 0.10 if r_done: bonus += 0.10 if n_done: bonus += 0.05 if contained: bonus += 0.10 if self._state.attacker_succeeded: bonus -= 0.20 if not (q_done or r_done or n_done): bonus -= 0.05 self._state.plan_progress["conclude"] = True self._state.total_reward += bonus self._state.incident_contained = True req = self._scenario["required_actions"] missing_packages = sorted(set(req["quarantine"]) - set(self._state.quarantined)) missing_secrets = sorted(set(req["rotate_secret"]) - set(self._state.rotated_secrets)) missing_teams = sorted(set(req["notify"]) - set(self._state.notified_teams)) result = ( f"{'OK' if q_done else 'MISS'} Packages quarantined: " f"{len(self._state.quarantined)}/{len(req['quarantine'])}" + (f" - missing: {missing_packages}" if missing_packages else "") + "\n" + f"{'OK' if r_done else 'MISS'} Secrets rotated: " + f"{len(self._state.rotated_secrets)}/{len(req['rotate_secret'])}" + (f" - missing: {missing_secrets}" if missing_secrets else "") + "\n" + f"{'OK' if n_done else 'MISS'} Teams notified: " + f"{len(self._state.notified_teams)}/{len(req['notify'])}" + (f" - missing: {missing_teams}" if missing_teams else "") + "\n" + ( "Contained before exfiltration" if contained else "Incident closed before full containment" ) ) return self._terminal_obs(bonus, result=result) def _invalid_target(self, label: str, value: str, steps_left: int) -> AgenticSecurityLabObservation: self._state.invalid_action_count += 1 reward = -0.01 self._state.total_reward += reward return self._build_obs( reward=reward, success=False, result=f"{label} '{value}' not found.", steps_left=steps_left, error=f"{label} not found", ) def _remember_package(self, package: str) -> None: if package not in self._state.discovered_packages: self._state.discovered_packages.append(package) def _discover_secrets_for_package(self, package: str) -> None: affected_owners = set(self._state.dependents.get(package, [])) for secret_name, secret_meta in self._state.secrets.items(): if secret_meta["owner"] in affected_owners and secret_name not in self._state.discovered_secrets: self._state.discovered_secrets.append(secret_name) def _score_breakdown(self) -> dict[str, float]: required = self._scenario["required_actions"] quarantine_ratio = self._ratio(self._state.quarantined, required["quarantine"]) rotate_ratio = self._ratio(self._state.rotated_secrets, required["rotate_secret"]) notify_ratio = self._ratio(self._state.notified_teams, required["notify"]) contain_ratio = 1.0 if (not self._state.attacker_succeeded and quarantine_ratio == 1.0) else 0.0 return { "quarantine_ratio": quarantine_ratio, "rotate_ratio": rotate_ratio, "notify_ratio": notify_ratio, "contain_ratio": contain_ratio, } def _benchmark_score(self) -> float: return benchmark_score_from_breakdown(self._score_breakdown()) @staticmethod def _ratio(actual: list[str], required: list[str]) -> float: if not required: return 1.0 return len(set(actual) & set(required)) / len(set(required)) def _build_obs( self, reward: float, success: bool, result: str, steps_left: int, data: dict[str, Any] | None = None, error: str | None = None, ) -> AgenticSecurityLabObservation: benchmark_score = self._benchmark_score() score_breakdown = self._score_breakdown() evaluator_metrics = { "invalid_actions": self._state.invalid_action_count, "false_positives": self._state.false_positive_count, "mode_fallback_used": self._state.mode_fallback_used, "mode": self._state.mode, "command_fallback_enabled": self._state.command_fallback_enabled, "command_fallback_used_count": self._state.command_fallback_used_count, "deadline_reached": self._state.step_count >= self._state.exfiltration_step, "attacker_succeeded": self._state.attacker_succeeded, } uncertainty = self._uncertainty_score() uncertainty_reduction = uncertainty_reduction_score(self._last_uncertainty, uncertainty) self._last_uncertainty = uncertainty observation_data = { "reward_type": "training_step_reward", "benchmark_score": benchmark_score, "score_breakdown": score_breakdown, "process_breakdown": { "uncertainty_reduction": uncertainty_reduction, }, "evaluator_metrics": evaluator_metrics, "max_steps": self._state.max_steps, "packages_in_scope": list(self._state.packages.keys()), } if data: observation_data.update(data) done = ( self._state.incident_contained or self._state.attacker_succeeded or self._state.step_count >= self._state.max_steps ) return AgenticSecurityLabObservation( success=success, done=done, reward=reward, result=result, data=observation_data, incident_summary=( f"Step {self._state.step_count}/{self._state.max_steps} | " f"KnownPackages:{len(self._state.discovered_packages)} " f"KnownSecrets:{len(self._state.discovered_secrets)} " f"Q:{len(self._state.quarantined)} " f"R:{len(self._state.rotated_secrets)} " f"N:{len(self._state.notified_teams)} " f"Attacker:{self._state.attacker_progress:.2f}" ), steps_remaining=steps_left, exposed_secrets=[ secret for secret in self._state.discovered_secrets if not self._state.secrets[secret]["rotated"] ], active_malicious_packages=[ package for package in self._state.discovered_packages if package not in self._state.quarantined ], visible_alerts=self._state.discovered_iocs[-5:], uncertainty_score=uncertainty, plan_progress=dict(self._state.plan_progress), info={ "attacker_progress": self._state.attacker_progress, "risk_events": self._state.risk_events[-3:], "discovered_packages": len(self._state.discovered_packages), "discovered_secrets": len(self._state.discovered_secrets), }, error=error, ) def _terminal_obs(self, reward: float, result: str = "") -> AgenticSecurityLabObservation: self._state.incident_contained = True return self._build_obs( reward=reward, success=True, result=result or "Episode ended.", steps_left=0, ) def _uncertainty_score(self) -> float: total_malicious = sum(1 for meta in self._state.packages.values() if meta["malicious"]) total_known = len(self._state.discovered_packages) unresolved = max(0, total_malicious - total_known) base = unresolved / max(1, total_malicious) hidden = len(self._state.pending_hidden_iocs) / max(1, len(self._scenario.get("hidden_iocs", []))) return round(max(0.0, min(1.0, 0.15 + 0.5 * base + 0.35 * hidden)), 4) def _advance_attacker(self) -> None: if self._state.attacker_succeeded: return if self._state.mode == "benchmark": delta = 1.0 / max(1, self._state.exfiltration_step) else: jitter = self._scenario.get("stochastic", {}).get("progress_jitter", 0.3) delta = max( 0.05, (1.0 / max(1, self._state.exfiltration_step)) + self._rng.uniform(-jitter, jitter) * 0.05, ) self._state.attacker_progress = min(1.0, self._state.attacker_progress + delta) if self._state.attacker_progress > 0.85 and not self._state.risk_events: self._state.risk_events.append("Attacker foothold appears to be deepening.") if self._state.step_count >= self._state.exfiltration_step: critical_remaining = [ secret_name for secret_name, secret_meta in self._state.secrets.items() if secret_meta["critical"] and not secret_meta["rotated"] ] if critical_remaining: self._state.attacker_succeeded = True self._state.risk_events.append("Critical secrets exfiltrated.") def _log_transition(self, command: str, params: dict[str, Any], obs: AgenticSecurityLabObservation) -> None: self._state.trajectory_log.append( { "step": self._state.step_count, "mode": self._state.mode, "command": command, "params": params, "reward": obs.reward, "done": obs.done, "success": obs.success, "benchmark_score": obs.data.get("benchmark_score", 0.0), "attacker_progress": self._state.attacker_progress, "plan_progress": dict(self._state.plan_progress), } )