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,131 @@
+"""MITRE ATT&CK mapper — maps extracted TTPs to ATT&CK technique IDs."""
+import json
+import hashlib
+import logging
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+ATTACK_DATA_PATH = Path(__file__).parent.parent / 'data' / 'mitre_attack.json'
+
+
+class MITREMapper:
+    """Map extracted TTPs to MITRE ATT&CK technique IDs."""
+
+    def __init__(self):
+        self.techniques = {}
+        self.tactics = {}
+        self._loaded = False
+        self._load_attack_data()
+
+    def _load_attack_data(self):
+        """Load ATT&CK data and verify integrity."""
+        if not ATTACK_DATA_PATH.exists():
+            logger.warning("MITRE ATT&CK data not found at %s", ATTACK_DATA_PATH)
+            return
+
+        try:
+            with open(ATTACK_DATA_PATH) as f:
+                data = json.load(f)
+
+            # Load tactics
+            for tactic in data.get('tactics', []):
+                tid = tactic['id']
+                self.tactics[tid] = tactic['name']
+                self.tactics[tactic['name'].lower()] = tid
+
+            # Load techniques (if populated)
+            for technique in data.get('techniques', []):
+                tid = technique.get('id', '')
+                name = technique.get('name', '')
+                if tid and name:
+                    self.techniques[tid] = {
+                        'name': name,
+                        'tactic': technique.get('tactic', ''),
+                        'description': technique.get('description', ''),
+                    }
+                    self.techniques[name.lower()] = tid
+
+            self._loaded = True
+            logger.debug("Loaded %d tactics, %d techniques",
+                        len(data.get('tactics', [])),
+                        len(data.get('techniques', [])))
+
+        except Exception as e:
+            logger.error("Failed to load ATT&CK data: %s", e)
+
+    def map_technique(self, technique_name: str, category: str = "") -> Optional[str]:
+        """Map a technique name or description to a MITRE ATT&CK ID.
+
+        Returns ATT&CK ID (e.g., 'T1566') or None if no match.
+        """
+        if not self._loaded:
+            return None
+
+        # Direct ID match
+        if technique_name.upper().startswith('T') and technique_name[1:].replace('.', '').isdigit():
+            if technique_name.upper() in self.techniques:
+                return technique_name.upper()
+
+        # Name match
+        lower_name = technique_name.lower().strip()
+        if lower_name in self.techniques:
+            result = self.techniques[lower_name]
+            if isinstance(result, str):
+                return result
+
+        # Category to tactic mapping
+        category_tactic_map = {
+            'initial_access': 'TA0001',
+            'execution': 'TA0002',
+            'persistence': 'TA0003',
+            'privilege_escalation': 'TA0004',
+            'defense_evasion': 'TA0005',
+            'credential_access': 'TA0006',
+            'discovery': 'TA0007',
+            'lateral_movement': 'TA0008',
+            'collection': 'TA0009',
+            'exfiltration': 'TA0010',
+            'command_and_control': 'TA0011',
+            'impact': 'TA0040',
+            'resource_development': 'TA0042',
+            'reconnaissance': 'TA0043',
+            # Aliases
+            'c2': 'TA0011',
+            'exfil': 'TA0010',
+            'privesc': 'TA0004',
+            'recon': 'TA0043',
+            'surveillance': 'TA0009',
+        }
+
+        return category_tactic_map.get(category.lower().strip())
+
+    def get_tactic_name(self, tactic_id: str) -> Optional[str]:
+        """Get tactic name from ID."""
+        return self.tactics.get(tactic_id)
+
+    def get_technique_info(self, technique_id: str) -> Optional[dict]:
+        """Get technique details from ID."""
+        return self.techniques.get(technique_id) if isinstance(
+            self.techniques.get(technique_id), dict) else None
+
+    def enrich_ttp(self, technique: str, category: str) -> dict:
+        """Enrich a TTP with MITRE mapping.
+
+        Returns dict with mitre_id and tactic if found.
+        """
+        result = {'mitre_id': None, 'tactic': None}
+
+        mitre_id = self.map_technique(technique, category)
+        if mitre_id:
+            if mitre_id.startswith('TA'):
+                result['tactic'] = mitre_id
+                result['tactic_name'] = self.get_tactic_name(mitre_id)
+            else:
+                result['mitre_id'] = mitre_id
+                info = self.get_technique_info(mitre_id)
+                if info:
+                    result['tactic'] = info.get('tactic')
+
+        return result