| 1 | import asyncio |
| 2 | import os |
| 3 | from datetime import datetime |
| 4 | from datetime import timedelta |
| 5 | from typing import Optional |
| 6 | |
| 7 | import httpx |
| 8 | from loguru import logger |
| 9 | |
| 10 | from app.integrations.copilot_searches.schema.copilot_searches import PlatformFilter |
| 11 | from app.integrations.copilot_searches.schema.copilot_searches import RuleSeverity |
| 12 | from app.integrations.copilot_searches.schema.copilot_searches import RuleStatus |
| 13 | from app.integrations.copilot_searches.services.copilot_searches import rules_cache |
| 14 | |
| 15 | MITRE_STIX_URL = "https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json" |
| 16 | MITRE_CACHE_TTL_HOURS = 24 |
| 17 | |
| 18 | |
| 19 | class MitreMatrix: |
| 20 | """In-memory cache of the MITRE ATT&CK Enterprise matrix structure. |
| 21 | |
| 22 | Pulls the official STIX bundle from mitre/cti and indexes tactics + techniques |
| 23 | so we can cross-reference them against CoPilot Search rules. |
| 24 | """ |
| 25 | |
| 26 | def __init__(self) -> None: |
| 27 | self._tactics: list[dict] = [] |
| 28 | self._techniques: dict[str, dict] = {} |
| 29 | self._last_refresh: Optional[datetime] = None |
| 30 | self._lock = asyncio.Lock() |
| 31 | |
| 32 | @property |
| 33 | def is_stale(self) -> bool: |
| 34 | if self._last_refresh is None: |
| 35 | return True |
| 36 | return datetime.utcnow() - self._last_refresh > timedelta(hours=MITRE_CACHE_TTL_HOURS) |
| 37 | |
| 38 | async def ensure_loaded(self) -> None: |
| 39 | if self.is_stale: |
| 40 | await self.refresh() |
| 41 | |
| 42 | async def refresh(self) -> None: |
| 43 | async with self._lock: |
| 44 | logger.info(f"Fetching MITRE ATT&CK Enterprise STIX bundle from {MITRE_STIX_URL}") |
| 45 | headers = {} |
| 46 | token = os.getenv("GITHUB_TOKEN") |
| 47 | if token: |
| 48 | headers["Authorization"] = f"Bearer {token}" |
| 49 | |
| 50 | async with httpx.AsyncClient(timeout=120.0, headers=headers) as client: |
| 51 | response = await client.get(MITRE_STIX_URL) |
| 52 | response.raise_for_status() |
| 53 | bundle = response.json() |
| 54 | |
| 55 | self._parse_bundle(bundle) |
| 56 | self._last_refresh = datetime.utcnow() |
| 57 | logger.info( |
| 58 | f"Loaded MITRE matrix: {len(self._tactics)} tactics, {len(self._techniques)} techniques", |
| 59 | ) |
| 60 | |
| 61 | def _parse_bundle(self, bundle: dict) -> None: |
| 62 | tactics_by_short: dict[str, dict] = {} |
| 63 | techniques: dict[str, dict] = {} |
| 64 | |
| 65 | for obj in bundle.get("objects", []): |
| 66 | obj_type = obj.get("type") |
| 67 | |
| 68 | if obj.get("revoked") or obj.get("x_mitre_deprecated"): |
| 69 | continue |
| 70 | |
| 71 | if obj_type == "x-mitre-tactic": |
| 72 | external_id = self._get_external_id(obj) |
| 73 | short_name = obj.get("x_mitre_shortname", "") |
| 74 | if not external_id or not short_name: |
| 75 | continue |
| 76 | tactics_by_short[short_name] = { |
| 77 | "id": external_id, |
| 78 | "name": obj.get("name", ""), |
| 79 | "short_name": short_name, |
| 80 | "url": self._get_external_url(obj), |
| 81 | "description": obj.get("description", ""), |
| 82 | } |
| 83 | |
| 84 | elif obj_type == "attack-pattern": |
| 85 | external_id = self._get_external_id(obj) |
| 86 | if not external_id: |
| 87 | continue |
| 88 | tactic_short_names = [ |
| 89 | phase.get("phase_name", "") |
| 90 | for phase in obj.get("kill_chain_phases", []) |
| 91 | if phase.get("kill_chain_name") == "mitre-attack" |
| 92 | ] |
| 93 | techniques[external_id] = { |
| 94 | "id": external_id, |
| 95 | "name": obj.get("name", ""), |
| 96 | "url": self._get_external_url(obj), |
| 97 | "is_subtechnique": bool(obj.get("x_mitre_is_subtechnique")), |
| 98 | "tactic_short_names": tactic_short_names, |
| 99 | } |
| 100 | |
| 101 | # Order tactics by the canonical kill-chain order. The STIX bundle stores |
| 102 | # the official ordering in the matrix object's tactic_refs; we approximate |
| 103 | # with the conventional order here as a fallback. |
| 104 | canonical_order = [ |
| 105 | "reconnaissance", |
| 106 | "resource-development", |
| 107 | "initial-access", |
| 108 | "execution", |
| 109 | "persistence", |
| 110 | "privilege-escalation", |
| 111 | "defense-evasion", |
| 112 | "credential-access", |
| 113 | "discovery", |
| 114 | "lateral-movement", |
| 115 | "collection", |
| 116 | "command-and-control", |
| 117 | "exfiltration", |
| 118 | "impact", |
| 119 | ] |
| 120 | ordered_tactics: list[dict] = [] |
| 121 | for short in canonical_order: |
| 122 | if short in tactics_by_short: |
| 123 | ordered_tactics.append(tactics_by_short[short]) |
| 124 | # Append any tactics not in the canonical list (forward-compat) |
| 125 | for short, t in tactics_by_short.items(): |
| 126 | if short not in canonical_order: |
| 127 | ordered_tactics.append(t) |
| 128 | |
| 129 | self._tactics = ordered_tactics |
| 130 | self._techniques = techniques |
| 131 | |
| 132 | @staticmethod |
| 133 | def _get_external_id(obj: dict) -> str: |
| 134 | for ref in obj.get("external_references", []): |
| 135 | if ref.get("source_name") == "mitre-attack": |
| 136 | return ref.get("external_id", "") |
| 137 | return "" |
| 138 | |
| 139 | @staticmethod |
| 140 | def _get_external_url(obj: dict) -> str: |
| 141 | for ref in obj.get("external_references", []): |
| 142 | if ref.get("source_name") == "mitre-attack": |
| 143 | return ref.get("url", "") |
| 144 | return "" |
| 145 | |
| 146 | @property |
| 147 | def tactics(self) -> list[dict]: |
| 148 | return self._tactics |
| 149 | |
| 150 | @property |
| 151 | def techniques(self) -> dict[str, dict]: |
| 152 | return self._techniques |
| 153 | |
| 154 | |
| 155 | mitre_matrix = MitreMatrix() |
| 156 | |
| 157 | |
| 158 | def _rule_matches_filters( |
| 159 | rule: dict, |
| 160 | platform: Optional[PlatformFilter], |
| 161 | severity: Optional[RuleSeverity], |
| 162 | status: Optional[RuleStatus], |
| 163 | has_graylog: Optional[bool], |
| 164 | search: Optional[str], |
| 165 | ) -> bool: |
| 166 | if platform is not None and platform != PlatformFilter.ALL: |
| 167 | if rule.get("_platform", "unknown") != platform.value: |
| 168 | return False |
| 169 | if severity is not None: |
| 170 | if rule.get("response", {}).get("severity", "").lower() != severity.value: |
| 171 | return False |
| 172 | if status is not None: |
| 173 | if rule.get("status", "").lower() != status.value: |
| 174 | return False |
| 175 | if has_graylog is not None: |
| 176 | if rule.get("_has_graylog", False) != has_graylog: |
| 177 | return False |
| 178 | if search: |
| 179 | s = search.lower() |
| 180 | name = rule.get("name", "").lower() |
| 181 | desc = rule.get("description", "").lower() |
| 182 | if s not in name and s not in desc: |
| 183 | return False |
| 184 | return True |
| 185 | |
| 186 | |
| 187 | async def get_coverage( |
| 188 | platform: Optional[PlatformFilter] = None, |
| 189 | severity: Optional[RuleSeverity] = None, |
| 190 | status: Optional[RuleStatus] = None, |
| 191 | has_graylog: Optional[bool] = None, |
| 192 | search: Optional[str] = None, |
| 193 | ) -> dict: |
| 194 | """Build the MITRE coverage map by cross-referencing rules against the matrix. |
| 195 | |
| 196 | Optional filters narrow the rules considered (platform/severity/status/has_graylog/search) |
| 197 | so the matrix can show "Windows-only coverage", etc. |
| 198 | |
| 199 | Returns a payload shaped for the frontend matrix view: ordered tactic columns, |
| 200 | techniques grouped under each tactic, per-technique rule counts + IDs with |
| 201 | sub-techniques nested, and a flat `rules_index` mapping rule ID to a small |
| 202 | summary (name, severity, platform, has_graylog) for hover previews. |
| 203 | """ |
| 204 | await mitre_matrix.ensure_loaded() |
| 205 | await rules_cache.ensure_loaded() |
| 206 | |
| 207 | # Map base technique -> {rule_ids set, subtechniques: {sub_id -> rule_ids set}} |
| 208 | coverage: dict[str, dict] = {} |
| 209 | rules_index: dict[str, dict] = {} |
| 210 | |
| 211 | for rule in rules_cache.get_all_rules(): |
| 212 | rule_id = rule.get("id", "") |
| 213 | if not rule_id: |
| 214 | continue |
| 215 | if not _rule_matches_filters(rule, platform, severity, status, has_graylog, search): |
| 216 | continue |
| 217 | |
| 218 | # Cap data_sources at 3 to keep payload small; full list is in /id/{rule_id}. |
| 219 | ds = rule.get("data_source", []) or [] |
| 220 | rules_index[rule_id] = { |
| 221 | "id": rule_id, |
| 222 | "name": rule.get("name", ""), |
| 223 | "severity": rule.get("response", {}).get("severity", "medium"), |
| 224 | "platform": rule.get("_platform", "unknown"), |
| 225 | "has_graylog": rule.get("_has_graylog", False), |
| 226 | "data_sources": [s for s in ds if isinstance(s, str)][:3], |
| 227 | } |
| 228 | |
| 229 | for raw_tid in rule.get("tags", {}).get("mitre_attack_id", []) or []: |
| 230 | tid = raw_tid.strip().upper() |
| 231 | if not tid.startswith("T"): |
| 232 | continue |
| 233 | base = tid.split(".")[0] |
| 234 | entry = coverage.setdefault(base, {"rule_ids": set(), "subtechniques": {}}) |
| 235 | if "." in tid: |
| 236 | sub = entry["subtechniques"].setdefault(tid, set()) |
| 237 | sub.add(rule_id) |
| 238 | else: |
| 239 | entry["rule_ids"].add(rule_id) |
| 240 | |
| 241 | # Build tactic-column structure |
| 242 | techniques_meta = mitre_matrix.techniques |
| 243 | tactics_out: list[dict] = [] |
| 244 | |
| 245 | for tactic in mitre_matrix.tactics: |
| 246 | techniques_in_tactic: list[dict] = [] |
| 247 | for tid, meta in techniques_meta.items(): |
| 248 | if meta["is_subtechnique"]: |
| 249 | continue |
| 250 | if tactic["short_name"] not in meta["tactic_short_names"]: |
| 251 | continue |
| 252 | |
| 253 | cov = coverage.get(tid, {"rule_ids": set(), "subtechniques": {}}) |
| 254 | base_rule_ids = sorted(cov["rule_ids"]) |
| 255 | sub_entries: list[dict] = [] |
| 256 | total_with_subs = set(cov["rule_ids"]) |
| 257 | |
| 258 | for sub_tid, sub_meta in techniques_meta.items(): |
| 259 | if not sub_meta["is_subtechnique"]: |
| 260 | continue |
| 261 | if not sub_tid.startswith(tid + "."): |
| 262 | continue |
| 263 | sub_rule_ids = sorted(cov["subtechniques"].get(sub_tid, set())) |
| 264 | total_with_subs.update(sub_rule_ids) |
| 265 | sub_entries.append( |
| 266 | { |
| 267 | "id": sub_tid, |
| 268 | "name": sub_meta["name"], |
| 269 | "url": sub_meta["url"], |
| 270 | "rule_count": len(sub_rule_ids), |
| 271 | "rule_ids": sub_rule_ids, |
| 272 | }, |
| 273 | ) |
| 274 | sub_entries.sort(key=lambda s: s["id"]) |
| 275 | |
| 276 | techniques_in_tactic.append( |
| 277 | { |
| 278 | "id": tid, |
| 279 | "name": meta["name"], |
| 280 | "url": meta["url"], |
| 281 | "rule_count": len(base_rule_ids), |
| 282 | "rule_ids": base_rule_ids, |
| 283 | "total_rule_count": len(total_with_subs), |
| 284 | "subtechniques": sub_entries, |
| 285 | }, |
| 286 | ) |
| 287 | |
| 288 | techniques_in_tactic.sort(key=lambda t: t["id"]) |
| 289 | |
| 290 | tactics_out.append( |
| 291 | { |
| 292 | "id": tactic["id"], |
| 293 | "name": tactic["name"], |
| 294 | "short_name": tactic["short_name"], |
| 295 | "url": tactic["url"], |
| 296 | "techniques": techniques_in_tactic, |
| 297 | }, |
| 298 | ) |
| 299 | |
| 300 | total_techniques = sum(len(t["techniques"]) for t in tactics_out) |
| 301 | covered_techniques = sum(1 for t in tactics_out for tech in t["techniques"] if tech["total_rule_count"] > 0) |
| 302 | |
| 303 | return { |
| 304 | "success": True, |
| 305 | "message": "MITRE coverage built successfully", |
| 306 | "tactics": tactics_out, |
| 307 | "rules_index": rules_index, |
| 308 | "stats": { |
| 309 | "total_tactics": len(tactics_out), |
| 310 | "total_techniques": total_techniques, |
| 311 | "covered_techniques": covered_techniques, |
| 312 | "total_rules": len(rules_index), |
| 313 | "matrix_last_refreshed": mitre_matrix._last_refresh, |
| 314 | "rules_last_refreshed": rules_cache.last_refresh, |
| 315 | }, |
| 316 | } |