File size: 1,818 Bytes
df5d3a2 b2aea8e df5d3a2 b2aea8e df5d3a2 b2aea8e 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 RemediationActivityEvent
def is_503_error(e):
return isinstance(e, APIError) and (getattr(e, 'code', None) == 503 or "503" in str(e))
class GreenTeamFixer:
def __init__(self):
self.name = "green-team-fixer"
# Load system prompt from green_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, "green_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:
"""Executes Remediation Activity based on the provided JSON string containing the finding and JIT token."""
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=RemediationActivityEvent,
),
)
# Validate output against OCSF schema
event_dict = json.loads(response.text)
event = RemediationActivityEvent(**event_dict)
return event.model_dump_json()
|