Add extractors: TTP, tool, tradecraft, surveillance, infrastructure, entity, MITRE mapper + prompts

2026-03-19 08:51:37 -04:00
parent 41cbdf5b04
commit ca53524118
10 changed files with 909 additions and 0 deletions
@@ -0,0 +1,51 @@
+"""TTP (Techniques, Tactics, Procedures) extractor for Mosaic."""
+import logging
+from typing import Optional
+
+from .base import BaseExtractor
+from storage.db import MosaicDB
+from storage.models import TTP
+from analysis.prompts import TTP_SYSTEM_PROMPT, TTP_EXTRACTION_PROMPT
+
+logger = logging.getLogger(__name__)
+
+
+class TTPExtractor(BaseExtractor):
+    EXTRACTOR_NAME = "ttp"
+
+    def get_system_prompt(self) -> str:
+        return TTP_SYSTEM_PROMPT
+
+    def get_extraction_prompt(self, text: str) -> str:
+        return TTP_EXTRACTION_PROMPT.format(document=text)
+
+    def get_json_schema(self) -> dict:
+        return {
+            "properties": {
+                "items": {"type": "array"},
+            },
+            "required": ["items"]
+        }
+
+    def process_response(self, data: dict, doc_id: str) -> int:
+        items = data.get('items', [])
+        if isinstance(data, list):
+            items = data
+        count = 0
+        for item in items:
+            try:
+                ttp = TTP(
+                    technique=item.get('technique', ''),
+                    category=item.get('category', ''),
+                    description=item.get('description', ''),
+                    actor=item.get('actor'),
+                    mitre_id=item.get('mitre_id'),
+                    source_doc_ids=[doc_id],
+                    source_context=item.get('source_context', '')[:500],
+                )
+                ttp.compute_dedup_key()
+                if self.db.insert_ttp(ttp):
+                    count += 1
+            except Exception as e:
+                logger.warning("Failed to process TTP item: %s", e)
+        return count