File size: 7,804 Bytes
df5d3a2 a1b1008 df5d3a2 a1b1008 df5d3a2 a1b1008 df5d3a2 16fe371 df5d3a2 16fe371 df5d3a2 a1b1008 df5d3a2 a1b1008 df5d3a2 a1b1008 df5d3a2 8beda30 df5d3a2 16fe371 8beda30 df5d3a2 16fe371 df5d3a2 8beda30 9101ed3 df5d3a2 a1b1008 df5d3a2 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import os
import json
from google import genai
from src.schemas.ocsf import NetworkActivityEvent, DetectionFindingEvent, RemediationActivityEvent
from src.agents.red_team.agent import RedTeamAttacker
from src.agents.blue_team.agent import BlueTeamDefender
from src.agents.green_team.agent import GreenTeamFixer
from google.adk.agents.base_agent import BaseAgent
from google.adk.events.event import Event
from google.genai.errors import APIError
from pydantic import PrivateAttr
from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception
from typing import AsyncGenerator
def is_503_error(e):
return isinstance(e, APIError) and (getattr(e, 'code', None) == 503 or "503" in str(e))
class OrchestratorCoordinator(BaseAgent):
name: str = "orchestrator-coordinator"
description: str = "Event-driven coordinator for the simulation workflow."
_system_prompt: str = PrivateAttr()
_client: genai.Client = PrivateAttr()
_red_team: RedTeamAttacker = PrivateAttr()
_blue_team: BlueTeamDefender = PrivateAttr()
_green_team: GreenTeamFixer = PrivateAttr()
def model_post_init(self, __context):
super().model_post_init(__context)
# Load system prompt from orchestrator.md in the project root
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
md_path = os.path.join(project_root, "orchestrator.md")
with open(md_path, "r", encoding="utf-8") as f:
self._system_prompt = f.read()
# client is instantiated dynamically in invoke()
self._red_team = RedTeamAttacker()
self._blue_team = BlueTeamDefender()
self._green_team = GreenTeamFixer()
@retry(
wait=wait_exponential(min=2, max=8),
stop=stop_after_attempt(3),
retry=retry_if_exception(is_503_error),
reraise=True
)
async def invoke(self, prompt: str, **kwargs) -> str:
"""Runs the 3-phase A2A simulation based on the prompt (target context)."""
api_key = kwargs.get("api_key")
simulation_report = {
"simulation_id": "sim-uuid",
"phases": {
"attack": {"ocsf_event": {}, "attack_vector": ""},
"evaluate": {"ocsf_finding": {}, "agent_trust_score": 1.0, "circuit_breaker_tripped": False},
"remediate": {"vibe_diff": "", "hitl_approved": False, "quarantine_status": "", "refactored_code": ""}
},
"jit_tokens_issued": [],
"simulation_outcome": "DETECTED"
}
# Phase 1: Attack (Red Team)
red_output = await self._red_team.invoke(prompt, api_key=api_key)
# Validate and store
red_event = NetworkActivityEvent.model_validate_json(red_output)
simulation_report["phases"]["attack"]["ocsf_event"] = red_event.model_dump()
if red_event.unmapped:
simulation_report["phases"]["attack"]["attack_vector"] = red_event.unmapped.attack_vector
# Phase 2: Evaluate (Blue Team)
blue_output = await self._blue_team.invoke(red_output, api_key=api_key)
# Validate and store
blue_finding = DetectionFindingEvent.model_validate_json(blue_output)
simulation_report["phases"]["evaluate"]["ocsf_finding"] = blue_finding.model_dump()
trust_score = 1.0
violations = 0
if blue_finding.unmapped:
# Recalculate trust score deterministically from ABA results.
# Never rely on the LLM's self-reported trust score — it hallucinates.
aba = blue_finding.unmapped.aba_check_results
violations = sum([
aba.agbom_violation,
aba.execution_loop_detected,
aba.prompt_injection_detected,
aba.semantic_drift_detected,
])
trust_score = round(max(0.0, 1.0 - violations * 0.3), 1)
# Trip circuit breaker on ANY detected violation — a SOC investigates everything.
circuit_breaker = violations > 0
simulation_report["phases"]["evaluate"]["agent_trust_score"] = trust_score
simulation_report["phases"]["evaluate"]["circuit_breaker_tripped"] = circuit_breaker
# Phase 3 — HITL gate: trigger on any ABA violation
if violations > 0:
# Generate Vibe Diff — use async API to avoid blocking the event loop
try:
vibe_diff_prompt = (
f"System: {self._system_prompt}\n"
f"Translate this Detection Finding into a plain-English Vibe Diff.\n"
f"List exactly three bullet sections:\n"
f"1. Revoke: what tool/access will be revoked and why\n"
f"2. Refactor: what code will be hardened and how\n"
f"3. Unchanged: what will NOT be changed\n"
f"Use backticks for code/tool names.\n\n"
f"Detection Finding:\n{blue_output}"
)
client = genai.Client(api_key=api_key)
vibe_diff_response = await client.aio.models.generate_content(
model='gemini-2.5-flash',
contents=vibe_diff_prompt
)
vibe_diff_text = vibe_diff_response.text or "Vibe diff generation returned empty response."
except Exception as vd_err:
vibe_diff_text = f"Vibe diff generation failed: {vd_err}"
simulation_report["phases"]["remediate"]["vibe_diff"] = vibe_diff_text
simulation_report["simulation_outcome"] = "PENDING_HITL"
# Simulated HITL approval handling for the playground demo
if "approve" in prompt.lower():
jit_token = {"id": "jit-1234", "expires_at": "2030-01-01T00:00:00Z", "allowed_actions": ["stateful_quarantine", "auto_refactoring"]}
simulation_report["jit_tokens_issued"].append(jit_token)
# Combine finding and token to send to green team
green_input = json.dumps({
"finding": blue_finding.model_dump(),
"jit_token": jit_token
})
green_output = await self._green_team.invoke(green_input, api_key=api_key)
green_event = RemediationActivityEvent.model_validate_json(green_output)
if green_event.unmapped:
simulation_report["phases"]["remediate"]["quarantine_status"] = green_event.unmapped.quarantine_status
simulation_report["phases"]["remediate"]["refactored_code"] = green_event.unmapped.refactored_code
simulation_report["simulation_outcome"] = "REMEDIATED"
return json.dumps(simulation_report, indent=2)
async def _run_async_impl(self, ctx) -> AsyncGenerator[Event, None]:
# Extract prompt from user_content or events
prompt = ""
if ctx.user_content and ctx.user_content.parts:
prompt = "".join(part.text for part in ctx.user_content.parts if part.text)
else:
events = ctx._get_events(current_invocation=True)
for event in reversed(events):
if event.author == "user" and event.content and event.content.parts:
prompt = "".join(part.text for part in event.content.parts if part.text)
break
response_str = await self.invoke(prompt)
yield Event(
invocation_id=ctx.invocation_id,
author=self.name,
branch=ctx.branch,
message=response_str
)
root_agent = OrchestratorCoordinator()
|