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

extractors/tool_extractor.py

@@ -0,0 +1,78 @@
+"""Named tool/exploit/implant extractor for Mosaic."""
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+
+from .base import BaseExtractor
+from storage.db import MosaicDB
+from storage.models import ExtractedTool
+from analysis.prompts import TOOL_SYSTEM_PROMPT, TOOL_EXTRACTION_PROMPT
+
+logger = logging.getLogger(__name__)
+
+KNOWN_TOOLS_PATH = Path(__file__).parent.parent / 'data' / 'known_tools.json'
+
+
+class ToolExtractor(BaseExtractor):
+    EXTRACTOR_NAME = "tool"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.known_tools = self._load_known_tools()
+
+    def _load_known_tools(self) -> list[str]:
+        try:
+            with open(KNOWN_TOOLS_PATH) as f:
+                data = json.load(f)
+            names = []
+            for tool in data.get('tools', []):
+                names.append(tool['name'].lower())
+                for alias in tool.get('aliases', []):
+                    names.append(alias.lower())
+            return names
+        except Exception as e:
+            logger.warning("Could not load known tools: %s", e)
+            return []
+
+    def get_system_prompt(self) -> str:
+        return TOOL_SYSTEM_PROMPT
+
+    def get_extraction_prompt(self, text: str) -> str:
+        # Include known tools as context
+        known_hint = ', '.join(self.known_tools[:50]) if self.known_tools else 'none loaded'
+        return TOOL_EXTRACTION_PROMPT.format(document=text, known_tools=known_hint)
+
+    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:
+                tool = ExtractedTool(
+                    name=item.get('name', ''),
+                    aliases=item.get('aliases', []),
+                    description=item.get('description', ''),
+                    capability=item.get('capability', ''),
+                    target_platforms=item.get('target_platforms', []),
+                    target_software=item.get('target_software', []),
+                    cves=item.get('cves', []),
+                    source_doc_ids=[doc_id],
+                    source_context=item.get('source_context', '')[:500],
+                    mitre_techniques=item.get('mitre_techniques', []),
+                )
+                tool.compute_dedup_key()
+                if self.db.insert_extracted_tool(tool):
+                    count += 1
+            except Exception as e:
+                logger.warning("Failed to process tool item: %s", e)
+        return count