mitre-attack-mcp 0.1.0__py3-none-any.whl

mitre_attack_mcp/data.py

@@ -0,0 +1,260 @@
+"""Download and indexing layer for MITRE ATT&CK STIX bundles."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Any, Iterable
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+DOMAINS: tuple[str, ...] = ("enterprise-attack", "ics-attack", "mobile-attack")
+
+BASE_URL = (
+    "https://raw.githubusercontent.com/mitre-attack/attack-stix-data/"
+    "master/{domain}/{domain}.json"
+)
+
+_DEFAULT_DATA_DIR = Path(__file__).resolve().parents[2] / "data"
+ENV_DATA_DIR = "MITRE_ATTACK_MCP_DATA_DIR"
+
+DOWNLOAD_TIMEOUT = 120.0
+DOWNLOAD_RETRIES = 3
+RETRY_BACKOFF = 2.0
+CHUNK_SIZE = 1 << 16
+
+
+# --------------------------------------------------------------------------- #
+# Paths
+# --------------------------------------------------------------------------- #
+
+
+def data_dir() -> Path:
+    override = os.environ.get(ENV_DATA_DIR)
+    return Path(override).expanduser() if override else _DEFAULT_DATA_DIR
+
+
+def file_for(domain: str) -> Path:
+    if domain not in DOMAINS:
+        raise ValueError(f"Unknown domain {domain!r}; expected one of {DOMAINS}")
+    return data_dir() / f"{domain}.json"
+
+
+# --------------------------------------------------------------------------- #
+# Download
+# --------------------------------------------------------------------------- #
+
+
+def download(domain: str, *, force: bool = False) -> Path:
+    if domain not in DOMAINS:
+        raise ValueError(f"Unknown domain {domain!r}; expected one of {DOMAINS}")
+
+    dest = file_for(domain)
+    if dest.exists() and dest.stat().st_size > 0 and not force:
+        logger.info("Using cached %s bundle at %s", domain, dest)
+        return dest
+
+    url = BASE_URL.format(domain=domain)
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+
+    last_exc: Exception | None = None
+    for attempt in range(1, DOWNLOAD_RETRIES + 1):
+        try:
+            logger.info(
+                "Downloading %s (attempt %d/%d)", domain, attempt, DOWNLOAD_RETRIES
+            )
+            _stream_to_file(url, tmp)
+            _validate_stix(tmp)
+            tmp.replace(dest)
+            logger.info(
+                "Saved %s bundle to %s (%d bytes)", domain, dest, dest.stat().st_size
+            )
+            return dest
+        except Exception as exc:  # noqa: BLE001
+            last_exc = exc
+            logger.warning("%s download attempt %d failed: %s", domain, attempt, exc)
+            tmp.unlink(missing_ok=True)
+            if attempt < DOWNLOAD_RETRIES:
+                time.sleep(RETRY_BACKOFF * attempt)
+
+    raise RuntimeError(
+        f"Failed to download {domain} bundle after {DOWNLOAD_RETRIES} attempts"
+    ) from last_exc
+
+
+def download_all(
+    domains: Iterable[str] | None = None,
+    *,
+    force: bool = False,
+) -> dict[str, Path]:
+    domains = tuple(domains or DOMAINS)
+    return {d: download(d, force=force) for d in domains}
+
+
+def _stream_to_file(url: str, dest: Path) -> None:
+    with httpx.stream("GET", url, timeout=DOWNLOAD_TIMEOUT, follow_redirects=True) as r:
+        r.raise_for_status()
+        total = int(r.headers.get("content-length", 0)) or None
+        written = 0
+        last_bucket = -1
+        with dest.open("wb") as fh:
+            for chunk in r.iter_bytes(CHUNK_SIZE):
+                fh.write(chunk)
+                written += len(chunk)
+                if total:
+                    bucket = (written * 100 // total) // 10
+                    if bucket != last_bucket:
+                        logger.debug("  ... %d%%", bucket * 10)
+                        last_bucket = bucket
+
+
+def _validate_stix(path: Path) -> None:
+    with path.open("r", encoding="utf-8") as fh:
+        payload = json.load(fh)
+    if not isinstance(payload, dict) or not isinstance(payload.get("objects"), list):
+        raise ValueError("Downloaded file is not a STIX bundle")
+
+
+# --------------------------------------------------------------------------- #
+# Indexing
+# --------------------------------------------------------------------------- #
+
+
+class AttackIndex:
+    """Indexed view over ONE domain's STIX bundle."""
+
+    def __init__(self, domain: str, objects: Iterable[dict[str, Any]]) -> None:
+        self.domain = domain
+        self.objects: list[dict[str, Any]] = list(objects)
+        self.by_attack_id: dict[str, dict[str, Any]] = {}
+        self.by_name: dict[str, dict[str, Any]] = {}
+        self.by_type: dict[str, list[dict[str, Any]]] = {}
+        self.aliases: dict[str, dict[str, Any]] = {}
+        self.by_stix_id: dict[str, dict[str, Any]] = {}
+        self._mitigates: dict[str, list[str]] = {}
+        for obj in self.objects:
+            self._index(obj)
+
+    def _index(self, obj: dict[str, Any]) -> None:
+        stix_type = obj.get("type")
+        if stix_type:
+            self.by_type.setdefault(stix_type, []).append(obj)
+
+        stix_id = obj.get("id")
+        if isinstance(stix_id, str):
+            self.by_stix_id[stix_id] = obj
+
+        name = obj.get("name")
+        if isinstance(name, str) and name:
+            self.by_name.setdefault(name.lower(), obj)
+
+        for alias in obj.get("aliases", []) or []:
+            if isinstance(alias, str) and alias:
+                self.aliases.setdefault(alias.lower(), obj)
+
+        for ref in obj.get("external_references", []) or []:
+            if ref.get("source_name") == "mitre-attack" and ref.get("external_id"):
+                self.by_attack_id[ref["external_id"].upper()] = obj
+
+        if stix_type == "relationship" and obj.get("relationship_type") == "mitigates":
+            src, tgt = obj.get("source_ref"), obj.get("target_ref")
+            if isinstance(src, str) and isinstance(tgt, str):
+                self._mitigates.setdefault(tgt, []).append(src)
+
+    def get_technique(self, attack_id: str) -> dict[str, Any] | None:
+        obj = self.by_attack_id.get(attack_id.strip().upper())
+        return obj if obj and obj.get("type") == "attack-pattern" else None
+
+    def get_group(self, name_or_alias: str) -> dict[str, Any] | None:
+        key = name_or_alias.strip().lower()
+        obj = self.by_name.get(key) or self.aliases.get(key)
+        return obj if obj and obj.get("type") == "intrusion-set" else None
+
+    def search_techniques(
+        self, keyword: str, *, limit: int = 10
+    ) -> list[dict[str, Any]]:
+        kw = keyword.strip().lower()
+        if not kw:
+            return []
+        hits: list[dict[str, Any]] = []
+        for obj in self.by_type.get("attack-pattern", []):
+            hay = f"{obj.get('name', '')} {obj.get('description', '')}".lower()
+            if kw in hay:
+                hits.append(obj)
+                if len(hits) >= limit:
+                    break
+        return hits
+
+    def mitigations_for_technique(self, technique_stix_id: str) -> list[dict[str, Any]]:
+        """Return mitigation objects linked to a technique via 'mitigates' relationships."""
+        source_ids = getattr(self, "_mitigates", {}).get(technique_stix_id, [])
+        results: list[dict[str, Any]] = []
+        for sid in source_ids:
+            obj = self.by_stix_id.get(sid)
+            if obj is None or obj.get("type") != "course-of-action":
+                continue
+            results.append(
+                {
+                    "id": attack_id_of(obj),
+                    "name": obj.get("name"),
+                    "description": (obj.get("description") or "")[:600],
+                }
+            )
+        return results
+
+    def counts(self) -> dict[str, int]:
+        return {
+            "techniques": len(self.by_type.get("attack-pattern", [])),
+            "groups": len(self.by_type.get("intrusion-set", [])),
+            "software": len(self.by_type.get("malware", []))
+            + len(self.by_type.get("tool", [])),
+            "mitigations": len(self.by_type.get("course-of-action", [])),
+        }
+
+
+_indices: dict[str, AttackIndex] = {}
+
+
+def get_index(domain: str) -> AttackIndex:
+    if domain not in DOMAINS:
+        raise ValueError(f"Unknown domain {domain!r}; expected one of {DOMAINS}")
+    if domain not in _indices:
+        path = file_for(domain)
+        if not path.exists():
+            raise FileNotFoundError(
+                f"ATT&CK data for {domain!r} not found at {path}. "
+                f"Run the server once with downloads enabled, or set "
+                f"{ENV_DATA_DIR} to a directory containing the bundle."
+            )
+        logger.info("Loading ATT&CK bundle for %s from %s", domain, path)
+        with path.open("r", encoding="utf-8") as fh:
+            bundle = json.load(fh)
+        _indices[domain] = AttackIndex(domain, bundle.get("objects", []))
+        logger.info("Indexed %s: %s", domain, _indices[domain].counts())
+    return _indices[domain]
+
+
+# --------------------------------------------------------------------------- #
+# STIX helpers
+# --------------------------------------------------------------------------- #
+
+
+def attack_id_of(obj: dict[str, Any]) -> str | None:
+    for ref in obj.get("external_references", []) or []:
+        if ref.get("source_name") == "mitre-attack" and ref.get("external_id"):
+            return ref["external_id"]
+    return None
+
+
+def tactics_of(technique: dict[str, Any]) -> list[str]:
+    return [
+        phase["phase_name"]
+        for phase in technique.get("kill_chain_phases", []) or []
+        if phase.get("kill_chain_name") == "mitre-attack"
+    ]

mitre_attack_mcp/server.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import logging
+import sys
+from mcp.server.mcpserver import MCPServer
+from . import data
+from . import services
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    stream=sys.stderr,  # never stdout — that's the JSON-RPC channel
+)
+
+mcp = MCPServer("mitre-attack-mcp", version="0.1.0")
+
+
+def _preStart():
+    logging.info("downloading mitre knowledge data..")
+    data.download_all()
+
+
+@mcp.tool()
+def get_data_stats(*, domain: str | None = None) -> dict | list[dict]:
+    """Return counts of techniques, groups, software, and mitigations.
+
+    domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
+        Omit to get a per-domain breakdown of all three.
+    """
+
+    return services.get_data_stats(domain=domain)
+
+
+@mcp.tool()
+def get_technique_by_id(
+    attack_id: str,
+    *,
+    domain: str | None = None,
+) -> dict:
+    """Look up a MITRE ATT&CK technique by its ID (e.g. 'T1059' or 'T1059.001').
+
+    Returns name, truncated description, tactics, and platforms, plus the
+    domain the match came from.
+
+    domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
+        Omit if you are not sure which domain the technique belongs to —
+        all three will be searched and the result will say which domain
+        matched. ATT&CK IDs are not globally unique across domains, so if
+        the ID exists in more than one domain the result also lists
+        `other_domains`.
+    """
+    return services.get_technique_by_id(attack_id, domain=domain)
+
+
+@mcp.tool()
+def search_techniques(
+    keyword: str,
+    *,
+    domain: str | None = None,
+    limit: int = 10,
+) -> list[dict]:
+    """Search techniques whose name or description contains the keyword.
+
+    domain: restrict to one of 'enterprise-attack', 'ics-attack',
+        'mobile-attack'. Omit to search all three domains.
+    limit: maximum results across all searched domains (default 10).
+    """
+    return services.search_techniques(keyword, domain=domain, limit=limit)
+
+
+@mcp.tool()
+def get_group_by_name(
+    name: str,
+    *,
+    domain: str | None = None,
+) -> dict:
+    """Look up a threat group by name or alias (e.g. 'APT29' or 'Cozy Bear').
+
+    domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
+        Omit if you are not sure which domain the group belongs to —
+        all three will be searched and the result will say which domain
+        matched.
+    """
+    return services.get_group_by_name(name, domain=domain)
+
+
+@mcp.tool()
+def get_mitigations_for_technique(
+    attack_id: str,
+    *,
+    domain: str | None = None,
+) -> dict:
+    """Return MITRE ATT&CK mitigations that apply to a specific technique.
+
+    Use this tool whenever the user asks how to defend against, mitigate,
+    or prevent a technique (e.g. "What mitigations exist for T1003?").
+    Returns each mitigation's ATT&CK ID (M####), name, and description,
+    sourced from the official ATT&CK dataset — do not answer from memory.
+
+    domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
+        Omit to search all three.
+    """
+    return services.get_mitigations_for_technique(attack_id, domain=domain)
+
+
+def main():
+    _preStart()
+    logging.info("mitre-attack-mcp server starting..")
+    mcp.run()
+    logging.info("mitre-attack-mcp server started..")
+
+
+if __name__ == "__main__":
+    main()

mitre_attack_mcp/services.py

@@ -0,0 +1,160 @@
+"""Pure business logic. This module must NOT import `mcp`."""
+
+from __future__ import annotations
+from . import data
+import logging
+import sys
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    stream=sys.stderr,  # never stdout — that's the JSON-RPC channel
+)
+
+
+def get_data_stats(*, domain: str | None = None) -> dict | list[dict]:
+    """Counts of techniques, groups, software, and mitigations.
+
+    domain: restrict to one domain, or omit for a per-domain breakdown.
+    """
+    if domain is not None:
+        idx = data.get_index(domain)
+        return {"domain": domain, **idx.counts()}
+
+    result = [{"domain": d, **data.get_index(d).counts()} for d in data.DOMAINS]
+    logging.info("results %s=", result)
+    return result
+
+
+def get_technique_by_id(
+    attack_id: str,
+    *,
+    domain: str | None = None,
+) -> dict:
+    """Look up a technique by ATT&CK ID.
+
+    If `domain` is omitted, all three domains are searched. The returned dict
+    includes `domain` so the caller knows where the match came from, and
+    `other_domains` lists any additional domains where the same ID also
+    exists (ATT&CK IDs are not globally unique across domains).
+    """
+    matches: list[tuple[str, dict]] = []
+    for d in (domain,) if domain else data.DOMAINS:
+        obj = data.get_index(d).get_technique(attack_id)
+        if obj is not None:
+            matches.append((d, obj))
+
+    if not matches:
+        where = domain or "any domain"
+        raise ValueError(f"Technique {attack_id!r} not found in {where}")
+
+    d, obj = matches[0]
+    result = {
+        "domain": d,
+        "id": data.attack_id_of(obj),
+        "name": obj.get("name"),
+        "description": (obj.get("description") or "")[:800],
+        "tactics": data.tactics_of(obj),
+        "platforms": obj.get("x_mitre_platforms", []),
+        "is_subtechnique": obj.get("x_mitre_is_subtechnique", False),
+    }
+    if len(matches) > 1:
+        result["other_domains"] = [m[0] for m in matches[1:]]
+        result["note"] = (
+            "This ID also exists in other domains with different content; "
+            "pass `domain` to disambiguate."
+        )
+    logging.info("results %s=", result)
+    return result
+
+
+def search_techniques(
+    keyword: str,
+    *,
+    domain: str | None = None,
+    limit: int = 10,
+) -> list[dict]:
+    """Search techniques by keyword. `domain=None` searches all domains."""
+    domains = (domain,) if domain else data.DOMAINS
+    for d in domains:
+        if d not in data.DOMAINS:
+            raise ValueError(f"Unknown domain {d!r}; expected one of {data.DOMAINS}")
+
+    results: list[dict] = []
+    for d in domains:
+        hits = data.get_index(d).search_techniques(keyword, limit=limit)
+        results.extend(
+            {"domain": d, "id": data.attack_id_of(o), "name": o.get("name")}
+            for o in hits
+        )
+        if len(results) >= limit:
+            break
+
+    logging.info("results %s=", results[:limit])
+    return results[:limit]
+
+
+def get_group_by_name(
+    name: str,
+    *,
+    domain: str | None = None,
+) -> dict:
+    """Look up a threat group by name or alias.
+
+    If `domain` is omitted, all three domains are searched; `domain` in the
+    result tells the caller where the match was found.
+    """
+    matches: list[tuple[str, dict]] = []
+    for d in (domain,) if domain else data.DOMAINS:
+        obj = data.get_index(d).get_group(name)
+        if obj is not None:
+            matches.append((d, obj))
+
+    if not matches:
+        where = domain or "any domain"
+        raise ValueError(f"Group {name!r} not found in {where}")
+
+    d, obj = matches[0]
+    result = {
+        "domain": d,
+        "name": obj.get("name"),
+        "aliases": obj.get("aliases", []),
+        "description": (obj.get("description") or "")[:800],
+    }
+    if len(matches) > 1:
+        result["other_domains"] = [m[0] for m in matches[1:]]
+
+    logging.info("results %s=", result)
+    return result
+
+
+def get_mitigations_for_technique(
+    attack_id: str,
+    *,
+    domain: str | None = None,
+) -> dict:
+    """Return mitigations linked to a technique.
+
+    If `domain` is omitted, all three domains are searched.
+    """
+    matches: list[tuple[str, dict]] = []
+    for d in (domain,) if domain else data.DOMAINS:
+        obj = data.get_index(d).get_technique(attack_id)
+        if obj is not None:
+            matches.append((d, obj))
+
+    if not matches:
+        where = domain or "any domain"
+        raise ValueError(f"Technique {attack_id!r} not found in {where}")
+
+    d, tech = matches[0]
+    mitigations = data.get_index(d).mitigations_for_technique(tech["id"])
+    result = {
+        "domain": d,
+        "technique_id": data.attack_id_of(tech),
+        "technique_name": tech.get("name"),
+        "mitigations": mitigations,
+    }
+    logging.info("results %s=", result)
+    return result

mitre_attack_mcp-0.1.0.dist-info/METADATA

@@ -0,0 +1,120 @@
+Metadata-Version: 2.5
+Name: mitre-attack-mcp
+Version: 0.1.0
+Summary: MITRE ATT&CK MCP server for AI agents
+Author: gsbadhan
+License: Apache-2.0
+Keywords: ai,attack,mcp,mitre,security,threat-intelligence
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Security
+Requires-Python: >=3.11
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: mcp[cli]>=2.3.0
+Description-Content-Type: text/markdown
+
+# mitre-attack-mcp
+MITRE ATT&CK MCP Server
+
+# Description
+A Model Context Protocol server that exposes the MITRE ATT&CK knowledge base to AI agents such as Claude Desktop and Claude Code.
+
+
+## What it does
+| Tool | Purpose | Input |
+|------|---------|-------|
+| `get_data_stats` | Counts of techniques, groups, software, mitigations | none |
+| `search_techniques` | Find techniques by keyword | `keyword` |
+| `get_technique_by_id` | Technique details by ID | `attack_id` |
+| `get_mitigations_for_technique` | Mitigations for a technique | `attack_id` |
+| `get_group_by_name` | Threat group details | `name` |
+
+All tools also accept an optional `domain`: `enterprise-attack`, `ics-attack`, or `mobile-attack`.
+
+
+# Project setup
+- uv venv
+- source .venv/bin/activate
+- uv sync
+
+# Data ingestion
+- MITRE data auto downloaded while MCP server starts
+
+# Run MCP server locally
+- uv run python -m mitre_attack_mcp.server
+#### Note: Dont run with Claude, its for local testing only
+
+# Unit tests
+- uv run pytest tests/ -v
+
+# Integration With Anthropic Claude Code
+```
+1. Clone code repository locally
+
+2. Check mcps: claude mcp list
+
+3. Add below configs in macOS: ~/Library/Application Support/Claude/claude_desktop_config.json
+
+  "mcpServers": {
+    "mitre-attack": {
+      "command": "uv",
+      "args": [
+        "run",
+        "--directory",
+        "/Users/john/git-projects/mitre-attack-mcp",
+        "mitre-attack-mcp"
+      ]
+    }
+  }
+
+
+4. Register mitre mcp: claude mcp add --scope user mitre-attack -- uv run --directory /Users/john/git-projects/mitre-attack-mcp mitre-attack-mcp
+
+5. Check mitre-attack mcp: claude mcp list
+
+6. Check detail: claude mcp get mitre-attack
+
+```
+# Integration with Anthropic Claude Desktop
+```
+1. Claude desktop -> Settings -> Developer -> Edit config
+
+2. Add below configs:
+
+  "mcpServers": {
+    "mitre-attack": {
+      "command": "uv",
+      "args": [
+        "run",
+        "--directory",
+        "/Users/john/git-projects/mitre-attack-mcp",
+        "mitre-attack-mcp"
+      ]
+    }
+  }
+
+3. Relaunch Claude
+
+```
+![Claude Settings](claude-settings.png)
+
+
+# Claude Desktop Assistant With Mitre MCP
+![Claude Desktop Chat](claude-desktop-chat.png)
+
+
+# Common prompts
+```
+Search mitre-attack for techniques containing 'credential'.
+Get the data stats from mitre-attack.
+Look up the group APT29 in mitre-attack.
+Use the mitre-attack server to look up technique T1059.
+what is mitigations for T1003.
+```
+
+

mitre_attack_mcp-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+mitre_attack_mcp/data.py,sha256=zK33D9zl7oQusBCS3bxROfkjpyTTWslnARIHwM8pKA4,9416
+mitre_attack_mcp/server.py,sha256=AzBlaLj9ivHq4cnJOjSn3KGozbUZHZ6D0aW1wDD0t7k,3447
+mitre_attack_mcp/services.py,sha256=pG8svJSkjd0qmkW4r6kbEYhHxMoKeHwzkeFR9Ji8Kq4,4923
+mitre_attack_mcp-0.1.0.dist-info/METADATA,sha256=JR-qq-g_BxPyl5m6f1CJOMRkUHd-xNE06XF2mgdOFpY,3156
+mitre_attack_mcp-0.1.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+mitre_attack_mcp-0.1.0.dist-info/entry_points.txt,sha256=S0AfUvIyWs4WXnsmdzt8tape431ydGkVUo-VrVwdL-U,66
+mitre_attack_mcp-0.1.0.dist-info/RECORD,,

mitre_attack_mcp-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

mitre_attack_mcp-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+mitre-attack-mcp = mitre_attack_mcp.server:main