File size: 1,801 Bytes
df5d3a2 b2aea8e df5d3a2 b2aea8e df5d3a2 b2aea8e 9101ed3 df5d3a2 9101ed3 | 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 | import os
import json
from google import genai
from google.genai import types
from google.genai.errors import APIError
from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception
from src.schemas.ocsf import DetectionFindingEvent
def is_503_error(e):
return isinstance(e, APIError) and (getattr(e, 'code', None) == 503 or "503" in str(e))
class BlueTeamDefender:
def __init__(self):
self.name = "blue-team-defender"
# Load system prompt from blue_team.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, "blue_team.md")
with open(md_path, "r", encoding="utf-8") as f:
self.system_prompt = f.read()
# client is instantiated dynamically in invoke()
@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, api_key: str = None) -> str:
"""Evaluates an OCSF 4001 event string and returns an OCSF 2004 Detection Finding JSON string."""
client = genai.Client(api_key=api_key)
response = client.models.generate_content(
model='gemini-2.5-flash',
contents=prompt,
config=types.GenerateContentConfig(
system_instruction=self.system_prompt,
response_mime_type="application/json",
response_schema=DetectionFindingEvent,
),
)
# Validate output against OCSF schema
event_dict = json.loads(response.text)
event = DetectionFindingEvent(**event_dict)
return event.model_dump_json()
|