n0mad1k
132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
"""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
|