Create app/agents/threat_modeling.py

app/agents/threat_modeling.py

@@ -0,0 +1,59 @@
+from typing import Dict, Any
+from app.utils.llm_client import LLMClient
+
+class BaseAgent:
+    def __init__(self, name: str, template: str, client: LLMClient):
+        self.name = name
+        self.template = template
+        self.client = client
+
+    def render(self, **kwargs) -> str:
+        return self.template.format(**kwargs)
+
+    def invoke(self, **kwargs) -> Any:
+        prompt = self.render(**kwargs)
+        return self.client.invoke(prompt)
+
+# --- Agent templates ---------------------------------------------------------
+
+THREAT_MODELER_PROMPT = """You are a cybersecurity threat-modeling expert.
+Input: {system_description}
+1. Identify assets and trust boundaries.
+2. Apply STRIDE categories.
+3. Map to MITRE ATT&CK techniques.
+Return JSON objects: asset, threat, stride_category, mitre_id, mitigation.
+"""
+
+THREAT_PATTERN_PROMPT = """Extract reusable threat patterns from this text.
+Input: {text_chunk}
+Return JSON list: pattern_name, description, preconditions, mitigations.
+"""
+
+ARCH_MAPPER_PROMPT = """Convert this architecture description into a DFD model.
+Input: {components_and_interactions}
+Return JSON: nodes, edges, trust_boundaries.
+"""
+
+# --- Agent classes -----------------------------------------------------------
+
+class ThreatModelerAgent(BaseAgent):
+    def run(self, system_description: str) -> Any:
+        return self.invoke(system_description=system_description)
+
+class ThreatPatternExtractor(BaseAgent):
+    def run(self, text_chunk: str) -> Any:
+        return self.invoke(text_chunk=text_chunk)
+
+class ArchitectureMapper(BaseAgent):
+    def run(self, components_and_interactions: str) -> Any:
+        return self.invoke(components_and_interactions=components_and_interactions)
+
+# --- Factory -----------------------------------------------------------------
+
+def build_agents(model_name: str = "gpt-threat-intel-v1"):
+    client = LLMClient(model_name)
+    return {
+        "modeler": ThreatModelerAgent("ThreatModelerAgent", THREAT_MODELER_PROMPT, client),
+        "extractor": ThreatPatternExtractor("ThreatPatternExtractor", THREAT_PATTERN_PROMPT, client),
+        "mapper": ArchitectureMapper("ArchitectureMapper", ARCH_MAPPER_PROMPT, client),
+    }