# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. """ Security Audit Environment Implementation. Simulates real-world VAPT engagements where an AI agent audits infrastructure for security vulnerabilities and compliance gaps. """ import random from copy import deepcopy from typing import Optional from uuid import uuid4 from openenv.core.env_server.interfaces import Environment try: from ..models import SecurityAuditAction, SecurityAuditObservation, SecurityAuditState except ImportError: from models import SecurityAuditAction, SecurityAuditObservation, SecurityAuditState try: from .scenarios import get_scenario, list_scenarios from .tools_engine import TOOL_DEFINITIONS, execute_tool from .grader import grade_episode, match_single_finding except ImportError: from server.scenarios import get_scenario, list_scenarios from server.tools_engine import TOOL_DEFINITIONS, execute_tool from server.grader import grade_episode, match_single_finding class SecurityAuditEnvironment(Environment): """ AI Security Audit Training Environment. Simulates real-world Vulnerability Assessment & Penetration Testing (VAPT) engagements. The agent discovers hosts, scans services, identifies vulnerabilities, and submits structured findings — just like a professional security auditor. Three scenarios with increasing difficulty: - Easy: Startup web app (2 hosts, 3 vulns) - Medium: E-commerce platform (4 hosts, 6 vulns) - Hard: Enterprise SOC2 audit (6 hosts, 10 vulns + honeypots) """ SUPPORTS_CONCURRENT_SESSIONS: bool = True # Difficulty multiplier for per-step tool/finding rewards _DIFFICULTY_REWARD_MULTIPLIER = {"easy": 1.0, "medium": 1.3, "hard": 1.6} def __init__(self): super().__init__() self._state = SecurityAuditState() self._scenario = None self._discovered_hosts: list = [] self._discovered_ports: dict = {} self._discovered_services: dict = {} self._submitted_findings: list = [] self._action_history: list = [] self._discovered_vulns: set = set() self._episode_reward: float = 0.0 self._last_tool_call: tuple = () self._rng: random.Random = random.Random() # Multi-agent extension: dynamic attack surface + sub-agent registry. # ``_attack_surface`` starts as the original scenario hosts and grows # whenever tools reveal new targets (e.g. SSRF dumping internal IPs). # ``_active_subagents`` tracks delegated investigation branches keyed # by spawn_id; each entry carries scope/target/budget/findings/parent_step. self._attack_surface: set = set() self._revealed_targets: list = [] self._active_subagents: dict = {} self._subagent_outcomes: list = [] # for grader's Delegation Score def reset(self, seed=None, episode_id=None, **kwargs) -> SecurityAuditObservation: """Reset the environment for a new audit engagement. kwargs: scenario_id: "easy", "medium", or "hard" (default: "easy") """ scenario_id = kwargs.get("scenario_id", "easy") self._scenario = deepcopy(get_scenario(scenario_id)) self._discovered_hosts = [] self._discovered_ports = {} self._discovered_services = {} self._submitted_findings = [] self._action_history = [] self._discovered_vulns = set() self._episode_reward = 0.0 self._last_tool_call = () self._rng = random.Random(seed) if seed is not None else random.Random() # Reset multi-agent state. Attack surface seeds from initial *visible* # hosts only — hidden hosts (those gated by `hidden_until`) only join # the attack surface when revealed by tool calls, which is what makes # spawn_subagent meaningful as a delegation primitive. self._attack_surface = { ip for ip, info in (self._scenario or {}).get("hosts", {}).items() if not info.get("hidden_until") } self._revealed_targets = [] self._active_subagents = {} self._subagent_outcomes = [] eid = episode_id or str(uuid4()) self._state = SecurityAuditState( episode_id=eid, step_count=0, scenario_id=scenario_id, scenario_name=self._scenario["name"], target_network=self._scenario["target_network"], max_steps=self._scenario["max_steps"], ) self._reset_rubric() return SecurityAuditObservation( tool_output="", message=self._scenario["briefing"], discovered_hosts=[], discovered_services={}, findings_submitted=0, steps_remaining=self._scenario["max_steps"], done=False, reward=0.0, ) def step(self, action: SecurityAuditAction, **kwargs) -> SecurityAuditObservation: self._state.step_count += 1 steps_remaining = self._state.max_steps - self._state.step_count # Sub-agent budget accounting: if a sub-agent is currently active, every # step (other than the spawn that created it, or its own return_to_parent) # consumes one budget unit. Auto-close on exhaustion so the env never # gets stuck in a sub-context. active_sid = self._active_subagent_id() if active_sid and action.action_type not in ("spawn_subagent",): sub = self._active_subagents[active_sid] sub["steps_used"] += 1 if sub["steps_used"] > sub["budget"] and action.action_type != "return_to_parent": # Auto-close as unproductive-by-timeout. Surface this to the # agent so it knows control is back at the parent. sub["status"] = "completed" sub["timeout"] = True findings_added = len(self._submitted_findings) - sub["findings_at_spawn"] productive = findings_added >= 1 sub["productive"] = productive self._subagent_outcomes.append({ "spawn_id": active_sid, "scope": sub["scope"], "target": sub["target"], "budget": sub["budget"], "steps_used": sub["steps_used"], "findings_added": findings_added, "productive": productive, "timeout": True, "spawn_step": sub["spawn_step"], "closed_at_step": self._state.step_count, }) # The agent's current action still runs at the parent layer. # We just clear the active sub-agent flag here. self._action_history.append({ "step": self._state.step_count, "action_type": action.action_type, "tool_name": action.tool_name, "arguments": action.arguments, "spawn_id": self._active_subagent_id(), }) if steps_remaining <= 0: return self._finish_episode("Step limit reached. Audit terminated.", truncated=True) if action.action_type == "list_tools": return self._handle_list_tools(steps_remaining) elif action.action_type == "use_tool": return self._handle_use_tool(action, steps_remaining) elif action.action_type == "submit_finding": return self._handle_submit_finding(action, steps_remaining) elif action.action_type == "spawn_subagent": return self._handle_spawn_subagent(action, steps_remaining) elif action.action_type == "return_to_parent": return self._handle_return_to_parent(action, steps_remaining) elif action.action_type == "generate_report": return self._finish_episode("Audit report generated.", truncated=False) else: return SecurityAuditObservation( tool_output=f"Unknown action_type: {action.action_type}", message="Use list_tools, use_tool, submit_finding, spawn_subagent, or generate_report.", discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=-0.05, ) def _active_subagent_id(self) -> Optional[str]: """Return the currently-active sub-agent spawn_id, if any.""" for sid, info in self._active_subagents.items(): if info.get("status") == "active": return sid return None def _handle_spawn_subagent(self, action: SecurityAuditAction, steps_remaining: int) -> SecurityAuditObservation: """Register a new sub-agent investigation branch. Args (in action.arguments): scope: "host" | "endpoint" | "cred" target: the IP / endpoint / credential identifier to investigate budget: integer step budget for the sub-agent (default 8, capped 15) Validates target is in the dynamic attack_surface (either originally scenario-discovered OR revealed by an earlier tool call). Registers the sub-agent and adds the target host to discovered_hosts so the sub-agent can immediately use scoped tools against it. """ if self._active_subagent_id() is not None: return self._obs_with_msg( "Cannot spawn — another sub-agent is still active. Call return_to_parent first.", steps_remaining, reward=-0.02, ) scope = (action.arguments or {}).get("scope", "host") target = (action.arguments or {}).get("target", "") budget = int((action.arguments or {}).get("budget", 8) or 8) budget = max(2, min(15, budget)) # clamp to a reasonable range if scope not in ("host", "endpoint", "cred"): return self._obs_with_msg( f"Invalid scope '{scope}'. Use one of: host | endpoint | cred.", steps_remaining, reward=-0.02, ) if not target: return self._obs_with_msg( "spawn_subagent requires a 'target' (e.g. 10.0.2.30 for scope=host).", steps_remaining, reward=-0.02, ) if scope == "host" and target not in self._attack_surface: return self._obs_with_msg( f"Target {target} not in current attack_surface. Discover or reveal it first.", steps_remaining, reward=-0.02, ) spawn_id = f"sub-{len(self._active_subagents) + 1:02d}-{self._state.step_count}" self._active_subagents[spawn_id] = { "scope": scope, "target": target, "budget": budget, "steps_used": 0, "findings_at_spawn": len(self._submitted_findings), "spawn_step": self._state.step_count, "parent_step": self._state.step_count, "status": "active", } # If host scope and the target is a revealed-but-not-yet-discovered host, # admit it to discovered_hosts now so sub-agent tools can hit it. if scope == "host" and target not in self._discovered_hosts and target in self._attack_surface: self._discovered_hosts.append(target) host_info = self._scenario.get("hosts", {}).get(target, {}) if self._scenario else {} if host_info.get("ports"): self._discovered_ports[target] = list(host_info.get("ports", [])) msg = ( f"Sub-agent {spawn_id} spawned: scope={scope} target={target} budget={budget} steps. " f"Subsequent actions are scoped to this branch until you call return_to_parent " f"(or budget exhausts)." ) return SecurityAuditObservation( tool_output=msg, message=msg, discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=0.01, # tiny positive — a valid spawn intent; productivity rewarded on return ) def _handle_return_to_parent(self, action: SecurityAuditAction, steps_remaining: int) -> SecurityAuditObservation: """Close the active sub-agent and record its outcome. Productivity: a sub-agent is "productive" if it submitted ≥1 finding during its run (we only count final-grader matches at episode end — for the per-step reward, any finding counts to give a fast signal). Reward: +0.05 productive -0.05 unproductive (penalises spurious branching) """ sid = self._active_subagent_id() if not sid: return self._obs_with_msg( "No active sub-agent to return from. Use spawn_subagent first.", steps_remaining, reward=-0.02, ) info = self._active_subagents[sid] findings_added = len(self._submitted_findings) - info["findings_at_spawn"] productive = findings_added >= 1 info["status"] = "completed" info["findings_added"] = findings_added info["productive"] = productive info["closed_at_step"] = self._state.step_count self._subagent_outcomes.append({ "spawn_id": sid, "scope": info["scope"], "target": info["target"], "budget": info["budget"], "steps_used": info["steps_used"], "findings_added": findings_added, "productive": productive, "spawn_step": info["spawn_step"], "closed_at_step": self._state.step_count, }) delegation_reward = 0.05 if productive else -0.05 msg = ( f"Sub-agent {sid} closed: {findings_added} findings submitted across " f"{info['steps_used']}/{info['budget']} steps. " f"{'PRODUCTIVE' if productive else 'unproductive'} (reward {delegation_reward:+.2f})." ) return SecurityAuditObservation( tool_output=msg, message=msg, discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=delegation_reward, ) def _obs_with_msg(self, msg: str, steps_remaining: int, reward: float = 0.0) -> SecurityAuditObservation: return SecurityAuditObservation( tool_output=msg, message=msg, discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=reward, ) @property def state(self) -> SecurityAuditState: self._state.discovered_hosts = list(self._discovered_hosts) self._state.discovered_ports = dict(self._discovered_ports) self._state.discovered_services = dict(self._discovered_services) self._state.submitted_findings = list(self._submitted_findings) self._state.total_reward = self._episode_reward return self._state def _current_phase(self) -> str: """Determine current audit phase from agent progress.""" if len(self._submitted_findings) > 0: return "exploitation" if len(self._discovered_hosts) > 0: return "enumeration" return "reconnaissance" # --- Action Handlers --- def _handle_list_tools(self, steps_remaining): tools_text = "Available security audit tools:\n\n" for tool in TOOL_DEFINITIONS: params = ", ".join(f"{k}: {v}" for k, v in tool["parameters"].items()) tools_text += f" {tool['name']}\n" tools_text += f" Description: {tool['description']}\n" tools_text += f" Parameters: {params}\n\n" return SecurityAuditObservation( tool_output=tools_text, available_tools=TOOL_DEFINITIONS, message="Use 'use_tool' action with tool_name and arguments to run a tool.", discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=0.0, ) def _handle_use_tool(self, action, steps_remaining): if not action.tool_name: return SecurityAuditObservation( tool_output="Error: tool_name is required for use_tool action.", message="Specify which tool to use.", discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=-0.02, ) target = action.arguments.get("host", action.arguments.get("target", "")) honeypots = self._scenario.get("honeypots", []) honeypot_penalty = -0.10 if target in honeypots else 0.0 # Detect redundant tool calls current_call = (action.tool_name, tuple(sorted(action.arguments.items()))) redundancy_penalty = -0.01 if current_call == self._last_tool_call else 0.0 self._last_tool_call = current_call output, new_hosts, new_ports, tool_reward, revealed = execute_tool( action.tool_name, action.arguments, self._scenario, self._discovered_hosts, self._discovered_ports, self._discovered_vulns, ) # Phase 1: any new revelations expand the dynamic attack_surface so # later spawn_subagent actions can validate their target. for r in revealed: tgt = r.get("target") if tgt and tgt not in self._attack_surface: self._attack_surface.add(tgt) self._revealed_targets.append({**r, "revealed_at_step": self._state.step_count}) # Difficulty multiplier on positive rewards difficulty = self._scenario.get("id", "easy") multiplier = self._DIFFICULTY_REWARD_MULTIPLIER.get(difficulty, 1.0) if tool_reward > 0: tool_reward *= multiplier for h in new_hosts: if h not in self._discovered_hosts: self._discovered_hosts.append(h) for h, ports in new_ports.items(): if h not in self._discovered_ports: self._discovered_ports[h] = [] for p in ports: if p not in self._discovered_ports[h]: self._discovered_ports[h].append(p) for h in self._discovered_hosts: host_ports = self._scenario.get("ports", {}).get(h, []) known_ports = self._discovered_ports.get(h, []) services = [f"{p['service']}/{p['port']}" for p in host_ports if p["port"] in known_ports] if services: self._discovered_services[h] = services reward = tool_reward + honeypot_penalty + redundancy_penalty self._episode_reward += reward message = "" if honeypot_penalty < 0: message = "WARNING: Activity detected on a monitoring system. This may be a honeypot." if redundancy_penalty < 0: message += " Note: Repeated identical tool call — consider a different action." return SecurityAuditObservation( tool_output=output, message=message.strip(), discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=reward, ) def _handle_submit_finding(self, action, steps_remaining): finding = action.arguments required = ["title", "host", "severity"] missing = [k for k in required if k not in finding] if missing: return SecurityAuditObservation( tool_output=f"Error: Missing required fields: {', '.join(missing)}", message="Finding must include at least: title, host, severity. " "Recommended: cvss_score, cwe, owasp, endpoint, evidence, remediation.", discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=-0.02, ) # Tag finding with current sub-agent context (if any) so the grader's # Delegation Score can credit sub-agent productivity correctly. active_sid = self._active_subagent_id() if active_sid: finding = {**finding, "_spawn_id": active_sid, "_parent_step": self._state.step_count} self._submitted_findings.append(finding) # Match using same logic as grader for consistency gt_vulns = self._scenario.get("vulnerabilities", []) matched_id = match_single_finding(finding, gt_vulns, self._discovered_vulns) difficulty = self._scenario.get("id", "easy") multiplier = self._DIFFICULTY_REWARD_MULTIPLIER.get(difficulty, 1.0) if matched_id: reward = 0.12 * multiplier self._discovered_vulns.add(matched_id) else: # Diminishing reward for unmatched findings to prevent spam unmatched = len(self._submitted_findings) - len(self._discovered_vulns) if unmatched <= 2: reward = 0.02 elif unmatched <= 4: reward = 0.01 else: reward = 0.0 self._episode_reward += reward return SecurityAuditObservation( tool_output=f"Finding #{len(self._submitted_findings)} recorded: {finding.get('title', 'Untitled')}", message=f"Finding submitted. Total findings: {len(self._submitted_findings)}.", discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=steps_remaining, current_phase=self._current_phase(), done=False, reward=reward, ) def _finish_episode(self, message, truncated=False): """End the audit and compute final grade.""" grades = grade_episode( self._scenario, self._submitted_findings, self._discovered_hosts, self._discovered_ports, steps_used=self._state.step_count, subagent_outcomes=self._subagent_outcomes, ) final_score = grades["final_score"] self._episode_reward += final_score report_lines = [ "=" * 60, "SECURITY AUDIT REPORT", "=" * 60, f"Scenario: {self._scenario['name']}", f"Company: {self._scenario['company']}", f"Compliance Framework: {self._scenario['compliance_context']}", "", "RESULTS:", f" Final Score: {final_score:.2f} / 1.00", f" Detection Rate: {grades['detection_rate']:.2f} ({grades['true_positives']}/{grades['total_vulnerabilities']} vulnerabilities found)", f" Coverage: {grades['coverage']:.2f} ({grades['hosts_examined']}/{grades['total_hosts']} hosts examined)", f" Severity Accuracy: {grades['severity_accuracy']:.2f}", f" Classification: CWE {grades['cwe_accuracy']:.2f} | OWASP {grades['owasp_accuracy']:.2f} | Combined {grades['classification_accuracy']:.2f}", f" Report Quality: {grades['report_quality']:.2f} (fields: {grades['field_completeness']:.2f}, narrative: {grades['narrative_quality']:.2f})", f" Pivoting Score: {grades['pivoting_score']:.2f}", f" Exploitation Proof: {grades['exploitation_proof']:.2f}", f" False Positives: {grades['false_positives']} (penalty: -{grades['fp_penalty']:.2f})", f" Honeypot Penalty: -{grades['honeypot_penalty']:.2f}", "", "COMPLIANCE:", f" Framework: {grades['compliance_framework']}", f" Controls Covered: {grades['compliance_controls_covered']}/{grades['compliance_controls_expected']}", f" Compliance Coverage: {grades['compliance_coverage']:.2f}", "", f"Steps Used: {self._state.step_count} / {self._scenario['max_steps']} (efficiency: {grades['efficiency']:.2f})", f"Findings Submitted: {len(self._submitted_findings)}", "=" * 60, ] return SecurityAuditObservation( tool_output="\n".join(report_lines), message=message, discovered_hosts=self._discovered_hosts, discovered_services=self._discovered_services, findings_submitted=len(self._submitted_findings), steps_remaining=0, done=True, truncated=truncated, current_phase="reporting", reward=final_score, metadata={"grades": grades}, )