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

extractors/infrastructure.py

@@ -0,0 +1,51 @@
+"""Infrastructure indicator extractor for Mosaic (IPs, domains, C2, hashes)."""
+import logging
+
+from .base import BaseExtractor
+from storage.models import Infrastructure
+from analysis.prompts import INFRASTRUCTURE_SYSTEM_PROMPT, INFRASTRUCTURE_EXTRACTION_PROMPT
+
+logger = logging.getLogger(__name__)
+
+
+class InfrastructureExtractor(BaseExtractor):
+    EXTRACTOR_NAME = "infrastructure"
+
+    def __init__(self, *args, collection: str = "", doc_title: str = "",
+                 doc_date: str = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.collection = collection
+        self.doc_title = doc_title
+        self.doc_date = doc_date
+
+    def get_system_prompt(self) -> str:
+        return INFRASTRUCTURE_SYSTEM_PROMPT
+
+    def get_extraction_prompt(self, text: str) -> str:
+        return INFRASTRUCTURE_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:
+                infra = Infrastructure(
+                    indicator=item.get('indicator', ''),
+                    indicator_type=item.get('indicator_type', ''),
+                    context=item.get('context', ''),
+                    source_doc_ids=[doc_id],
+                    source_context=item.get('source_context', '')[:500],
+                    collection=self.collection,
+                    doc_title=self.doc_title,
+                    doc_date=self.doc_date,
+                )
+                self.db.insert_infrastructure(infra)
+                count += 1
+            except Exception as e:
+                logger.warning("Failed to process infrastructure item: %s", e)
+        return count