mitre-attack-mcp 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of mitre-attack-mcp might be problematic. Click here for more details.
- mitre_attack_mcp-0.1.0/.gitignore +226 -0
- mitre_attack_mcp-0.1.0/PKG-INFO +120 -0
- mitre_attack_mcp-0.1.0/README.md +100 -0
- mitre_attack_mcp-0.1.0/claude-desktop-chat.png +0 -0
- mitre_attack_mcp-0.1.0/claude-settings.png +0 -0
- mitre_attack_mcp-0.1.0/pyproject.toml +39 -0
- mitre_attack_mcp-0.1.0/src/mitre_attack_mcp/data.py +260 -0
- mitre_attack_mcp-0.1.0/src/mitre_attack_mcp/server.py +115 -0
- mitre_attack_mcp-0.1.0/src/mitre_attack_mcp/services.py +160 -0
- mitre_attack_mcp-0.1.0/tests/test_server.py +38 -0
|
@@ -0,0 +1,226 @@
|
|
|
1
|
+
# avoid to commit mitre attack json files
|
|
2
|
+
data/enterprise-attack.json
|
|
3
|
+
data/ics-attack.json
|
|
4
|
+
data/mobile-attack.json
|
|
5
|
+
|
|
6
|
+
# Byte-compiled / optimized / DLL files
|
|
7
|
+
__pycache__/
|
|
8
|
+
*.py[codz]
|
|
9
|
+
*$py.class
|
|
10
|
+
|
|
11
|
+
# C extensions
|
|
12
|
+
*.so
|
|
13
|
+
|
|
14
|
+
# Distribution / packaging
|
|
15
|
+
.Python
|
|
16
|
+
build/
|
|
17
|
+
develop-eggs/
|
|
18
|
+
dist/
|
|
19
|
+
downloads/
|
|
20
|
+
eggs/
|
|
21
|
+
.eggs/
|
|
22
|
+
lib/
|
|
23
|
+
lib64/
|
|
24
|
+
parts/
|
|
25
|
+
sdist/
|
|
26
|
+
var/
|
|
27
|
+
wheels/
|
|
28
|
+
share/python-wheels/
|
|
29
|
+
*.egg-info/
|
|
30
|
+
.installed.cfg
|
|
31
|
+
*.egg
|
|
32
|
+
MANIFEST
|
|
33
|
+
|
|
34
|
+
# PyInstaller
|
|
35
|
+
# Usually these files are written by a python script from a template
|
|
36
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
37
|
+
*.manifest
|
|
38
|
+
*.spec
|
|
39
|
+
|
|
40
|
+
# Installer logs
|
|
41
|
+
pip-log.txt
|
|
42
|
+
pip-delete-this-directory.txt
|
|
43
|
+
|
|
44
|
+
# Unit test / coverage reports
|
|
45
|
+
htmlcov/
|
|
46
|
+
.tox/
|
|
47
|
+
.nox/
|
|
48
|
+
.coverage
|
|
49
|
+
.coverage.*
|
|
50
|
+
.cache
|
|
51
|
+
nosetests.xml
|
|
52
|
+
coverage.xml
|
|
53
|
+
*.cover
|
|
54
|
+
*.py.cover
|
|
55
|
+
*.lcov
|
|
56
|
+
.hypothesis/
|
|
57
|
+
.pytest_cache/
|
|
58
|
+
cover/
|
|
59
|
+
|
|
60
|
+
# Translations
|
|
61
|
+
*.mo
|
|
62
|
+
*.pot
|
|
63
|
+
|
|
64
|
+
# Django stuff:
|
|
65
|
+
*.log
|
|
66
|
+
local_settings.py
|
|
67
|
+
db.sqlite3
|
|
68
|
+
db.sqlite3-journal
|
|
69
|
+
|
|
70
|
+
# Flask stuff:
|
|
71
|
+
instance/
|
|
72
|
+
.webassets-cache
|
|
73
|
+
|
|
74
|
+
# Scrapy stuff:
|
|
75
|
+
.scrapy
|
|
76
|
+
|
|
77
|
+
# Sphinx documentation
|
|
78
|
+
docs/_build/
|
|
79
|
+
|
|
80
|
+
# PyBuilder
|
|
81
|
+
.pybuilder/
|
|
82
|
+
target/
|
|
83
|
+
|
|
84
|
+
# Jupyter Notebook
|
|
85
|
+
.ipynb_checkpoints
|
|
86
|
+
|
|
87
|
+
# IPython
|
|
88
|
+
profile_default/
|
|
89
|
+
ipython_config.py
|
|
90
|
+
|
|
91
|
+
# pyenv
|
|
92
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
93
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
94
|
+
# .python-version
|
|
95
|
+
|
|
96
|
+
# pipenv
|
|
97
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
98
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
99
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
100
|
+
# install all needed dependencies.
|
|
101
|
+
# Pipfile.lock
|
|
102
|
+
|
|
103
|
+
# uv
|
|
104
|
+
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
|
105
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
106
|
+
# commonly ignored for libraries.
|
|
107
|
+
# uv.lock
|
|
108
|
+
|
|
109
|
+
# poetry
|
|
110
|
+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
|
111
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
112
|
+
# commonly ignored for libraries.
|
|
113
|
+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
|
114
|
+
# poetry.lock
|
|
115
|
+
# poetry.toml
|
|
116
|
+
|
|
117
|
+
# pdm
|
|
118
|
+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
|
119
|
+
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
|
|
120
|
+
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
|
|
121
|
+
# pdm.lock
|
|
122
|
+
# pdm.toml
|
|
123
|
+
.pdm-python
|
|
124
|
+
.pdm-build/
|
|
125
|
+
|
|
126
|
+
# pixi
|
|
127
|
+
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
|
|
128
|
+
# pixi.lock
|
|
129
|
+
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
|
|
130
|
+
# in the .venv directory. It is recommended not to include this directory in version control.
|
|
131
|
+
.pixi/*
|
|
132
|
+
!.pixi/config.toml
|
|
133
|
+
|
|
134
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
135
|
+
__pypackages__/
|
|
136
|
+
|
|
137
|
+
# Celery stuff
|
|
138
|
+
celerybeat-schedule*
|
|
139
|
+
celerybeat.pid
|
|
140
|
+
|
|
141
|
+
# Redis
|
|
142
|
+
*.rdb
|
|
143
|
+
*.aof
|
|
144
|
+
*.pid
|
|
145
|
+
|
|
146
|
+
# RabbitMQ
|
|
147
|
+
mnesia/
|
|
148
|
+
rabbitmq/
|
|
149
|
+
rabbitmq-data/
|
|
150
|
+
|
|
151
|
+
# ActiveMQ
|
|
152
|
+
activemq-data/
|
|
153
|
+
|
|
154
|
+
# SageMath parsed files
|
|
155
|
+
*.sage.py
|
|
156
|
+
|
|
157
|
+
# Environments
|
|
158
|
+
.env
|
|
159
|
+
.envrc
|
|
160
|
+
.venv
|
|
161
|
+
uv.lock
|
|
162
|
+
env/
|
|
163
|
+
venv/
|
|
164
|
+
ENV/
|
|
165
|
+
env.bak/
|
|
166
|
+
venv.bak/
|
|
167
|
+
|
|
168
|
+
# Spyder project settings
|
|
169
|
+
.spyderproject
|
|
170
|
+
.spyproject
|
|
171
|
+
|
|
172
|
+
# Rope project settings
|
|
173
|
+
.ropeproject
|
|
174
|
+
|
|
175
|
+
# mkdocs/Zensical documentation
|
|
176
|
+
/site
|
|
177
|
+
|
|
178
|
+
# mypy
|
|
179
|
+
.mypy_cache/
|
|
180
|
+
.dmypy.json
|
|
181
|
+
dmypy.json
|
|
182
|
+
|
|
183
|
+
# Pyre type checker
|
|
184
|
+
.pyre/
|
|
185
|
+
|
|
186
|
+
# pytype static type analyzer
|
|
187
|
+
.pytype/
|
|
188
|
+
|
|
189
|
+
# Cython debug symbols
|
|
190
|
+
cython_debug/
|
|
191
|
+
|
|
192
|
+
# PyCharm
|
|
193
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
194
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
195
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
196
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
197
|
+
# .idea/
|
|
198
|
+
|
|
199
|
+
# Abstra
|
|
200
|
+
# Abstra is an AI-powered process automation framework.
|
|
201
|
+
# Ignore directories containing user credentials, local state, and settings.
|
|
202
|
+
# Learn more at https://abstra.io/docs
|
|
203
|
+
.abstra/
|
|
204
|
+
|
|
205
|
+
# Visual Studio Code
|
|
206
|
+
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore that
|
|
207
|
+
# can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
|
208
|
+
# and can be added to the global gitignore or merged into this file. However, if you prefer, you
|
|
209
|
+
# could uncomment the following to ignore the entire vscode folder
|
|
210
|
+
# .vscode/
|
|
211
|
+
# Temporary file for partial code execution
|
|
212
|
+
tempCodeRunnerFile.py
|
|
213
|
+
|
|
214
|
+
# Ruff stuff:
|
|
215
|
+
.ruff_cache/
|
|
216
|
+
|
|
217
|
+
# PyPI configuration file
|
|
218
|
+
.pypirc
|
|
219
|
+
|
|
220
|
+
# Marimo
|
|
221
|
+
marimo/_static/
|
|
222
|
+
marimo/_lsp/
|
|
223
|
+
__marimo__/
|
|
224
|
+
|
|
225
|
+
# Streamlit
|
|
226
|
+
.streamlit/secrets.toml
|
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: mitre-attack-mcp
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: MITRE ATT&CK MCP server for AI agents
|
|
5
|
+
Author: gsbadhan
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Keywords: ai,attack,mcp,mitre,security,threat-intelligence
|
|
8
|
+
Classifier: Development Status :: 3 - Alpha
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: Intended Audience :: Information Technology
|
|
11
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Topic :: Security
|
|
16
|
+
Requires-Python: >=3.11
|
|
17
|
+
Requires-Dist: httpx>=0.28.1
|
|
18
|
+
Requires-Dist: mcp[cli]>=2.3.0
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
|
|
21
|
+
# mitre-attack-mcp
|
|
22
|
+
MITRE ATT&CK MCP Server
|
|
23
|
+
|
|
24
|
+
# Description
|
|
25
|
+
A Model Context Protocol server that exposes the MITRE ATT&CK knowledge base to AI agents such as Claude Desktop and Claude Code.
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
## What it does
|
|
29
|
+
| Tool | Purpose | Input |
|
|
30
|
+
|------|---------|-------|
|
|
31
|
+
| `get_data_stats` | Counts of techniques, groups, software, mitigations | none |
|
|
32
|
+
| `search_techniques` | Find techniques by keyword | `keyword` |
|
|
33
|
+
| `get_technique_by_id` | Technique details by ID | `attack_id` |
|
|
34
|
+
| `get_mitigations_for_technique` | Mitigations for a technique | `attack_id` |
|
|
35
|
+
| `get_group_by_name` | Threat group details | `name` |
|
|
36
|
+
|
|
37
|
+
All tools also accept an optional `domain`: `enterprise-attack`, `ics-attack`, or `mobile-attack`.
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
# Project setup
|
|
41
|
+
- uv venv
|
|
42
|
+
- source .venv/bin/activate
|
|
43
|
+
- uv sync
|
|
44
|
+
|
|
45
|
+
# Data ingestion
|
|
46
|
+
- MITRE data auto downloaded while MCP server starts
|
|
47
|
+
|
|
48
|
+
# Run MCP server locally
|
|
49
|
+
- uv run python -m mitre_attack_mcp.server
|
|
50
|
+
#### Note: Dont run with Claude, its for local testing only
|
|
51
|
+
|
|
52
|
+
# Unit tests
|
|
53
|
+
- uv run pytest tests/ -v
|
|
54
|
+
|
|
55
|
+
# Integration With Anthropic Claude Code
|
|
56
|
+
```
|
|
57
|
+
1. Clone code repository locally
|
|
58
|
+
|
|
59
|
+
2. Check mcps: claude mcp list
|
|
60
|
+
|
|
61
|
+
3. Add below configs in macOS: ~/Library/Application Support/Claude/claude_desktop_config.json
|
|
62
|
+
|
|
63
|
+
"mcpServers": {
|
|
64
|
+
"mitre-attack": {
|
|
65
|
+
"command": "uv",
|
|
66
|
+
"args": [
|
|
67
|
+
"run",
|
|
68
|
+
"--directory",
|
|
69
|
+
"/Users/john/git-projects/mitre-attack-mcp",
|
|
70
|
+
"mitre-attack-mcp"
|
|
71
|
+
]
|
|
72
|
+
}
|
|
73
|
+
}
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
4. Register mitre mcp: claude mcp add --scope user mitre-attack -- uv run --directory /Users/john/git-projects/mitre-attack-mcp mitre-attack-mcp
|
|
77
|
+
|
|
78
|
+
5. Check mitre-attack mcp: claude mcp list
|
|
79
|
+
|
|
80
|
+
6. Check detail: claude mcp get mitre-attack
|
|
81
|
+
|
|
82
|
+
```
|
|
83
|
+
# Integration with Anthropic Claude Desktop
|
|
84
|
+
```
|
|
85
|
+
1. Claude desktop -> Settings -> Developer -> Edit config
|
|
86
|
+
|
|
87
|
+
2. Add below configs:
|
|
88
|
+
|
|
89
|
+
"mcpServers": {
|
|
90
|
+
"mitre-attack": {
|
|
91
|
+
"command": "uv",
|
|
92
|
+
"args": [
|
|
93
|
+
"run",
|
|
94
|
+
"--directory",
|
|
95
|
+
"/Users/john/git-projects/mitre-attack-mcp",
|
|
96
|
+
"mitre-attack-mcp"
|
|
97
|
+
]
|
|
98
|
+
}
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
3. Relaunch Claude
|
|
102
|
+
|
|
103
|
+
```
|
|
104
|
+

|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
# Claude Desktop Assistant With Mitre MCP
|
|
108
|
+

|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
# Common prompts
|
|
112
|
+
```
|
|
113
|
+
Search mitre-attack for techniques containing 'credential'.
|
|
114
|
+
Get the data stats from mitre-attack.
|
|
115
|
+
Look up the group APT29 in mitre-attack.
|
|
116
|
+
Use the mitre-attack server to look up technique T1059.
|
|
117
|
+
what is mitigations for T1003.
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
# mitre-attack-mcp
|
|
2
|
+
MITRE ATT&CK MCP Server
|
|
3
|
+
|
|
4
|
+
# Description
|
|
5
|
+
A Model Context Protocol server that exposes the MITRE ATT&CK knowledge base to AI agents such as Claude Desktop and Claude Code.
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
## What it does
|
|
9
|
+
| Tool | Purpose | Input |
|
|
10
|
+
|------|---------|-------|
|
|
11
|
+
| `get_data_stats` | Counts of techniques, groups, software, mitigations | none |
|
|
12
|
+
| `search_techniques` | Find techniques by keyword | `keyword` |
|
|
13
|
+
| `get_technique_by_id` | Technique details by ID | `attack_id` |
|
|
14
|
+
| `get_mitigations_for_technique` | Mitigations for a technique | `attack_id` |
|
|
15
|
+
| `get_group_by_name` | Threat group details | `name` |
|
|
16
|
+
|
|
17
|
+
All tools also accept an optional `domain`: `enterprise-attack`, `ics-attack`, or `mobile-attack`.
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
# Project setup
|
|
21
|
+
- uv venv
|
|
22
|
+
- source .venv/bin/activate
|
|
23
|
+
- uv sync
|
|
24
|
+
|
|
25
|
+
# Data ingestion
|
|
26
|
+
- MITRE data auto downloaded while MCP server starts
|
|
27
|
+
|
|
28
|
+
# Run MCP server locally
|
|
29
|
+
- uv run python -m mitre_attack_mcp.server
|
|
30
|
+
#### Note: Dont run with Claude, its for local testing only
|
|
31
|
+
|
|
32
|
+
# Unit tests
|
|
33
|
+
- uv run pytest tests/ -v
|
|
34
|
+
|
|
35
|
+
# Integration With Anthropic Claude Code
|
|
36
|
+
```
|
|
37
|
+
1. Clone code repository locally
|
|
38
|
+
|
|
39
|
+
2. Check mcps: claude mcp list
|
|
40
|
+
|
|
41
|
+
3. Add below configs in macOS: ~/Library/Application Support/Claude/claude_desktop_config.json
|
|
42
|
+
|
|
43
|
+
"mcpServers": {
|
|
44
|
+
"mitre-attack": {
|
|
45
|
+
"command": "uv",
|
|
46
|
+
"args": [
|
|
47
|
+
"run",
|
|
48
|
+
"--directory",
|
|
49
|
+
"/Users/john/git-projects/mitre-attack-mcp",
|
|
50
|
+
"mitre-attack-mcp"
|
|
51
|
+
]
|
|
52
|
+
}
|
|
53
|
+
}
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
4. Register mitre mcp: claude mcp add --scope user mitre-attack -- uv run --directory /Users/john/git-projects/mitre-attack-mcp mitre-attack-mcp
|
|
57
|
+
|
|
58
|
+
5. Check mitre-attack mcp: claude mcp list
|
|
59
|
+
|
|
60
|
+
6. Check detail: claude mcp get mitre-attack
|
|
61
|
+
|
|
62
|
+
```
|
|
63
|
+
# Integration with Anthropic Claude Desktop
|
|
64
|
+
```
|
|
65
|
+
1. Claude desktop -> Settings -> Developer -> Edit config
|
|
66
|
+
|
|
67
|
+
2. Add below configs:
|
|
68
|
+
|
|
69
|
+
"mcpServers": {
|
|
70
|
+
"mitre-attack": {
|
|
71
|
+
"command": "uv",
|
|
72
|
+
"args": [
|
|
73
|
+
"run",
|
|
74
|
+
"--directory",
|
|
75
|
+
"/Users/john/git-projects/mitre-attack-mcp",
|
|
76
|
+
"mitre-attack-mcp"
|
|
77
|
+
]
|
|
78
|
+
}
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
3. Relaunch Claude
|
|
82
|
+
|
|
83
|
+
```
|
|
84
|
+

|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
# Claude Desktop Assistant With Mitre MCP
|
|
88
|
+

|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
# Common prompts
|
|
92
|
+
```
|
|
93
|
+
Search mitre-attack for techniques containing 'credential'.
|
|
94
|
+
Get the data stats from mitre-attack.
|
|
95
|
+
Look up the group APT29 in mitre-attack.
|
|
96
|
+
Use the mitre-attack server to look up technique T1059.
|
|
97
|
+
what is mitigations for T1003.
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
|
|
Binary file
|
|
Binary file
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "mitre-attack-mcp"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "MITRE ATT&CK MCP server for AI agents"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.11"
|
|
7
|
+
license = { text = "Apache-2.0" }
|
|
8
|
+
authors = [{ name = "gsbadhan" }]
|
|
9
|
+
keywords = ["mcp", "mitre", "attack", "threat-intelligence", "security", "ai"]
|
|
10
|
+
classifiers = [
|
|
11
|
+
"Development Status :: 3 - Alpha",
|
|
12
|
+
"Intended Audience :: Developers",
|
|
13
|
+
"Intended Audience :: Information Technology",
|
|
14
|
+
"License :: OSI Approved :: Apache Software License",
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Programming Language :: Python :: 3.11",
|
|
17
|
+
"Programming Language :: Python :: 3.12",
|
|
18
|
+
"Topic :: Security",
|
|
19
|
+
]
|
|
20
|
+
dependencies = [
|
|
21
|
+
"httpx>=0.28.1",
|
|
22
|
+
"mcp[cli]>=2.3.0",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[project.scripts]
|
|
26
|
+
mitre-attack-mcp = "mitre_attack_mcp.server:main"
|
|
27
|
+
|
|
28
|
+
[build-system]
|
|
29
|
+
requires = ["hatchling"]
|
|
30
|
+
build-backend = "hatchling.build"
|
|
31
|
+
|
|
32
|
+
[tool.hatch.build.targets.wheel]
|
|
33
|
+
packages = ["src/mitre_attack_mcp"]
|
|
34
|
+
|
|
35
|
+
[dependency-groups]
|
|
36
|
+
dev = [
|
|
37
|
+
"pytest>=9.1.1",
|
|
38
|
+
"pytest-mcptest>=0.1.1",
|
|
39
|
+
]
|
|
@@ -0,0 +1,260 @@
|
|
|
1
|
+
"""Download and indexing layer for MITRE ATT&CK STIX bundles."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import logging
|
|
7
|
+
import os
|
|
8
|
+
import time
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from typing import Any, Iterable
|
|
11
|
+
|
|
12
|
+
import httpx
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger(__name__)
|
|
15
|
+
|
|
16
|
+
DOMAINS: tuple[str, ...] = ("enterprise-attack", "ics-attack", "mobile-attack")
|
|
17
|
+
|
|
18
|
+
BASE_URL = (
|
|
19
|
+
"https://raw.githubusercontent.com/mitre-attack/attack-stix-data/"
|
|
20
|
+
"master/{domain}/{domain}.json"
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
_DEFAULT_DATA_DIR = Path(__file__).resolve().parents[2] / "data"
|
|
24
|
+
ENV_DATA_DIR = "MITRE_ATTACK_MCP_DATA_DIR"
|
|
25
|
+
|
|
26
|
+
DOWNLOAD_TIMEOUT = 120.0
|
|
27
|
+
DOWNLOAD_RETRIES = 3
|
|
28
|
+
RETRY_BACKOFF = 2.0
|
|
29
|
+
CHUNK_SIZE = 1 << 16
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
# --------------------------------------------------------------------------- #
|
|
33
|
+
# Paths
|
|
34
|
+
# --------------------------------------------------------------------------- #
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def data_dir() -> Path:
|
|
38
|
+
override = os.environ.get(ENV_DATA_DIR)
|
|
39
|
+
return Path(override).expanduser() if override else _DEFAULT_DATA_DIR
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def file_for(domain: str) -> Path:
|
|
43
|
+
if domain not in DOMAINS:
|
|
44
|
+
raise ValueError(f"Unknown domain {domain!r}; expected one of {DOMAINS}")
|
|
45
|
+
return data_dir() / f"{domain}.json"
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
# --------------------------------------------------------------------------- #
|
|
49
|
+
# Download
|
|
50
|
+
# --------------------------------------------------------------------------- #
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def download(domain: str, *, force: bool = False) -> Path:
|
|
54
|
+
if domain not in DOMAINS:
|
|
55
|
+
raise ValueError(f"Unknown domain {domain!r}; expected one of {DOMAINS}")
|
|
56
|
+
|
|
57
|
+
dest = file_for(domain)
|
|
58
|
+
if dest.exists() and dest.stat().st_size > 0 and not force:
|
|
59
|
+
logger.info("Using cached %s bundle at %s", domain, dest)
|
|
60
|
+
return dest
|
|
61
|
+
|
|
62
|
+
url = BASE_URL.format(domain=domain)
|
|
63
|
+
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
64
|
+
tmp = dest.with_suffix(dest.suffix + ".part")
|
|
65
|
+
|
|
66
|
+
last_exc: Exception | None = None
|
|
67
|
+
for attempt in range(1, DOWNLOAD_RETRIES + 1):
|
|
68
|
+
try:
|
|
69
|
+
logger.info(
|
|
70
|
+
"Downloading %s (attempt %d/%d)", domain, attempt, DOWNLOAD_RETRIES
|
|
71
|
+
)
|
|
72
|
+
_stream_to_file(url, tmp)
|
|
73
|
+
_validate_stix(tmp)
|
|
74
|
+
tmp.replace(dest)
|
|
75
|
+
logger.info(
|
|
76
|
+
"Saved %s bundle to %s (%d bytes)", domain, dest, dest.stat().st_size
|
|
77
|
+
)
|
|
78
|
+
return dest
|
|
79
|
+
except Exception as exc: # noqa: BLE001
|
|
80
|
+
last_exc = exc
|
|
81
|
+
logger.warning("%s download attempt %d failed: %s", domain, attempt, exc)
|
|
82
|
+
tmp.unlink(missing_ok=True)
|
|
83
|
+
if attempt < DOWNLOAD_RETRIES:
|
|
84
|
+
time.sleep(RETRY_BACKOFF * attempt)
|
|
85
|
+
|
|
86
|
+
raise RuntimeError(
|
|
87
|
+
f"Failed to download {domain} bundle after {DOWNLOAD_RETRIES} attempts"
|
|
88
|
+
) from last_exc
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def download_all(
|
|
92
|
+
domains: Iterable[str] | None = None,
|
|
93
|
+
*,
|
|
94
|
+
force: bool = False,
|
|
95
|
+
) -> dict[str, Path]:
|
|
96
|
+
domains = tuple(domains or DOMAINS)
|
|
97
|
+
return {d: download(d, force=force) for d in domains}
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _stream_to_file(url: str, dest: Path) -> None:
|
|
101
|
+
with httpx.stream("GET", url, timeout=DOWNLOAD_TIMEOUT, follow_redirects=True) as r:
|
|
102
|
+
r.raise_for_status()
|
|
103
|
+
total = int(r.headers.get("content-length", 0)) or None
|
|
104
|
+
written = 0
|
|
105
|
+
last_bucket = -1
|
|
106
|
+
with dest.open("wb") as fh:
|
|
107
|
+
for chunk in r.iter_bytes(CHUNK_SIZE):
|
|
108
|
+
fh.write(chunk)
|
|
109
|
+
written += len(chunk)
|
|
110
|
+
if total:
|
|
111
|
+
bucket = (written * 100 // total) // 10
|
|
112
|
+
if bucket != last_bucket:
|
|
113
|
+
logger.debug(" ... %d%%", bucket * 10)
|
|
114
|
+
last_bucket = bucket
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def _validate_stix(path: Path) -> None:
|
|
118
|
+
with path.open("r", encoding="utf-8") as fh:
|
|
119
|
+
payload = json.load(fh)
|
|
120
|
+
if not isinstance(payload, dict) or not isinstance(payload.get("objects"), list):
|
|
121
|
+
raise ValueError("Downloaded file is not a STIX bundle")
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
# --------------------------------------------------------------------------- #
|
|
125
|
+
# Indexing
|
|
126
|
+
# --------------------------------------------------------------------------- #
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
class AttackIndex:
|
|
130
|
+
"""Indexed view over ONE domain's STIX bundle."""
|
|
131
|
+
|
|
132
|
+
def __init__(self, domain: str, objects: Iterable[dict[str, Any]]) -> None:
|
|
133
|
+
self.domain = domain
|
|
134
|
+
self.objects: list[dict[str, Any]] = list(objects)
|
|
135
|
+
self.by_attack_id: dict[str, dict[str, Any]] = {}
|
|
136
|
+
self.by_name: dict[str, dict[str, Any]] = {}
|
|
137
|
+
self.by_type: dict[str, list[dict[str, Any]]] = {}
|
|
138
|
+
self.aliases: dict[str, dict[str, Any]] = {}
|
|
139
|
+
self.by_stix_id: dict[str, dict[str, Any]] = {}
|
|
140
|
+
self._mitigates: dict[str, list[str]] = {}
|
|
141
|
+
for obj in self.objects:
|
|
142
|
+
self._index(obj)
|
|
143
|
+
|
|
144
|
+
def _index(self, obj: dict[str, Any]) -> None:
|
|
145
|
+
stix_type = obj.get("type")
|
|
146
|
+
if stix_type:
|
|
147
|
+
self.by_type.setdefault(stix_type, []).append(obj)
|
|
148
|
+
|
|
149
|
+
stix_id = obj.get("id")
|
|
150
|
+
if isinstance(stix_id, str):
|
|
151
|
+
self.by_stix_id[stix_id] = obj
|
|
152
|
+
|
|
153
|
+
name = obj.get("name")
|
|
154
|
+
if isinstance(name, str) and name:
|
|
155
|
+
self.by_name.setdefault(name.lower(), obj)
|
|
156
|
+
|
|
157
|
+
for alias in obj.get("aliases", []) or []:
|
|
158
|
+
if isinstance(alias, str) and alias:
|
|
159
|
+
self.aliases.setdefault(alias.lower(), obj)
|
|
160
|
+
|
|
161
|
+
for ref in obj.get("external_references", []) or []:
|
|
162
|
+
if ref.get("source_name") == "mitre-attack" and ref.get("external_id"):
|
|
163
|
+
self.by_attack_id[ref["external_id"].upper()] = obj
|
|
164
|
+
|
|
165
|
+
if stix_type == "relationship" and obj.get("relationship_type") == "mitigates":
|
|
166
|
+
src, tgt = obj.get("source_ref"), obj.get("target_ref")
|
|
167
|
+
if isinstance(src, str) and isinstance(tgt, str):
|
|
168
|
+
self._mitigates.setdefault(tgt, []).append(src)
|
|
169
|
+
|
|
170
|
+
def get_technique(self, attack_id: str) -> dict[str, Any] | None:
|
|
171
|
+
obj = self.by_attack_id.get(attack_id.strip().upper())
|
|
172
|
+
return obj if obj and obj.get("type") == "attack-pattern" else None
|
|
173
|
+
|
|
174
|
+
def get_group(self, name_or_alias: str) -> dict[str, Any] | None:
|
|
175
|
+
key = name_or_alias.strip().lower()
|
|
176
|
+
obj = self.by_name.get(key) or self.aliases.get(key)
|
|
177
|
+
return obj if obj and obj.get("type") == "intrusion-set" else None
|
|
178
|
+
|
|
179
|
+
def search_techniques(
|
|
180
|
+
self, keyword: str, *, limit: int = 10
|
|
181
|
+
) -> list[dict[str, Any]]:
|
|
182
|
+
kw = keyword.strip().lower()
|
|
183
|
+
if not kw:
|
|
184
|
+
return []
|
|
185
|
+
hits: list[dict[str, Any]] = []
|
|
186
|
+
for obj in self.by_type.get("attack-pattern", []):
|
|
187
|
+
hay = f"{obj.get('name', '')} {obj.get('description', '')}".lower()
|
|
188
|
+
if kw in hay:
|
|
189
|
+
hits.append(obj)
|
|
190
|
+
if len(hits) >= limit:
|
|
191
|
+
break
|
|
192
|
+
return hits
|
|
193
|
+
|
|
194
|
+
def mitigations_for_technique(self, technique_stix_id: str) -> list[dict[str, Any]]:
|
|
195
|
+
"""Return mitigation objects linked to a technique via 'mitigates' relationships."""
|
|
196
|
+
source_ids = getattr(self, "_mitigates", {}).get(technique_stix_id, [])
|
|
197
|
+
results: list[dict[str, Any]] = []
|
|
198
|
+
for sid in source_ids:
|
|
199
|
+
obj = self.by_stix_id.get(sid)
|
|
200
|
+
if obj is None or obj.get("type") != "course-of-action":
|
|
201
|
+
continue
|
|
202
|
+
results.append(
|
|
203
|
+
{
|
|
204
|
+
"id": attack_id_of(obj),
|
|
205
|
+
"name": obj.get("name"),
|
|
206
|
+
"description": (obj.get("description") or "")[:600],
|
|
207
|
+
}
|
|
208
|
+
)
|
|
209
|
+
return results
|
|
210
|
+
|
|
211
|
+
def counts(self) -> dict[str, int]:
|
|
212
|
+
return {
|
|
213
|
+
"techniques": len(self.by_type.get("attack-pattern", [])),
|
|
214
|
+
"groups": len(self.by_type.get("intrusion-set", [])),
|
|
215
|
+
"software": len(self.by_type.get("malware", []))
|
|
216
|
+
+ len(self.by_type.get("tool", [])),
|
|
217
|
+
"mitigations": len(self.by_type.get("course-of-action", [])),
|
|
218
|
+
}
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
_indices: dict[str, AttackIndex] = {}
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def get_index(domain: str) -> AttackIndex:
|
|
225
|
+
if domain not in DOMAINS:
|
|
226
|
+
raise ValueError(f"Unknown domain {domain!r}; expected one of {DOMAINS}")
|
|
227
|
+
if domain not in _indices:
|
|
228
|
+
path = file_for(domain)
|
|
229
|
+
if not path.exists():
|
|
230
|
+
raise FileNotFoundError(
|
|
231
|
+
f"ATT&CK data for {domain!r} not found at {path}. "
|
|
232
|
+
f"Run the server once with downloads enabled, or set "
|
|
233
|
+
f"{ENV_DATA_DIR} to a directory containing the bundle."
|
|
234
|
+
)
|
|
235
|
+
logger.info("Loading ATT&CK bundle for %s from %s", domain, path)
|
|
236
|
+
with path.open("r", encoding="utf-8") as fh:
|
|
237
|
+
bundle = json.load(fh)
|
|
238
|
+
_indices[domain] = AttackIndex(domain, bundle.get("objects", []))
|
|
239
|
+
logger.info("Indexed %s: %s", domain, _indices[domain].counts())
|
|
240
|
+
return _indices[domain]
|
|
241
|
+
|
|
242
|
+
|
|
243
|
+
# --------------------------------------------------------------------------- #
|
|
244
|
+
# STIX helpers
|
|
245
|
+
# --------------------------------------------------------------------------- #
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
def attack_id_of(obj: dict[str, Any]) -> str | None:
|
|
249
|
+
for ref in obj.get("external_references", []) or []:
|
|
250
|
+
if ref.get("source_name") == "mitre-attack" and ref.get("external_id"):
|
|
251
|
+
return ref["external_id"]
|
|
252
|
+
return None
|
|
253
|
+
|
|
254
|
+
|
|
255
|
+
def tactics_of(technique: dict[str, Any]) -> list[str]:
|
|
256
|
+
return [
|
|
257
|
+
phase["phase_name"]
|
|
258
|
+
for phase in technique.get("kill_chain_phases", []) or []
|
|
259
|
+
if phase.get("kill_chain_name") == "mitre-attack"
|
|
260
|
+
]
|
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
import sys
|
|
5
|
+
from mcp.server.mcpserver import MCPServer
|
|
6
|
+
from . import data
|
|
7
|
+
from . import services
|
|
8
|
+
|
|
9
|
+
logger = logging.getLogger(__name__)
|
|
10
|
+
logging.basicConfig(
|
|
11
|
+
level=logging.INFO,
|
|
12
|
+
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
|
13
|
+
stream=sys.stderr, # never stdout — that's the JSON-RPC channel
|
|
14
|
+
)
|
|
15
|
+
|
|
16
|
+
mcp = MCPServer("mitre-attack-mcp", version="0.1.0")
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def _preStart():
|
|
20
|
+
logging.info("downloading mitre knowledge data..")
|
|
21
|
+
data.download_all()
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@mcp.tool()
|
|
25
|
+
def get_data_stats(*, domain: str | None = None) -> dict | list[dict]:
|
|
26
|
+
"""Return counts of techniques, groups, software, and mitigations.
|
|
27
|
+
|
|
28
|
+
domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
|
|
29
|
+
Omit to get a per-domain breakdown of all three.
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
return services.get_data_stats(domain=domain)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@mcp.tool()
|
|
36
|
+
def get_technique_by_id(
|
|
37
|
+
attack_id: str,
|
|
38
|
+
*,
|
|
39
|
+
domain: str | None = None,
|
|
40
|
+
) -> dict:
|
|
41
|
+
"""Look up a MITRE ATT&CK technique by its ID (e.g. 'T1059' or 'T1059.001').
|
|
42
|
+
|
|
43
|
+
Returns name, truncated description, tactics, and platforms, plus the
|
|
44
|
+
domain the match came from.
|
|
45
|
+
|
|
46
|
+
domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
|
|
47
|
+
Omit if you are not sure which domain the technique belongs to —
|
|
48
|
+
all three will be searched and the result will say which domain
|
|
49
|
+
matched. ATT&CK IDs are not globally unique across domains, so if
|
|
50
|
+
the ID exists in more than one domain the result also lists
|
|
51
|
+
`other_domains`.
|
|
52
|
+
"""
|
|
53
|
+
return services.get_technique_by_id(attack_id, domain=domain)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@mcp.tool()
|
|
57
|
+
def search_techniques(
|
|
58
|
+
keyword: str,
|
|
59
|
+
*,
|
|
60
|
+
domain: str | None = None,
|
|
61
|
+
limit: int = 10,
|
|
62
|
+
) -> list[dict]:
|
|
63
|
+
"""Search techniques whose name or description contains the keyword.
|
|
64
|
+
|
|
65
|
+
domain: restrict to one of 'enterprise-attack', 'ics-attack',
|
|
66
|
+
'mobile-attack'. Omit to search all three domains.
|
|
67
|
+
limit: maximum results across all searched domains (default 10).
|
|
68
|
+
"""
|
|
69
|
+
return services.search_techniques(keyword, domain=domain, limit=limit)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
@mcp.tool()
|
|
73
|
+
def get_group_by_name(
|
|
74
|
+
name: str,
|
|
75
|
+
*,
|
|
76
|
+
domain: str | None = None,
|
|
77
|
+
) -> dict:
|
|
78
|
+
"""Look up a threat group by name or alias (e.g. 'APT29' or 'Cozy Bear').
|
|
79
|
+
|
|
80
|
+
domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
|
|
81
|
+
Omit if you are not sure which domain the group belongs to —
|
|
82
|
+
all three will be searched and the result will say which domain
|
|
83
|
+
matched.
|
|
84
|
+
"""
|
|
85
|
+
return services.get_group_by_name(name, domain=domain)
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
@mcp.tool()
|
|
89
|
+
def get_mitigations_for_technique(
|
|
90
|
+
attack_id: str,
|
|
91
|
+
*,
|
|
92
|
+
domain: str | None = None,
|
|
93
|
+
) -> dict:
|
|
94
|
+
"""Return MITRE ATT&CK mitigations that apply to a specific technique.
|
|
95
|
+
|
|
96
|
+
Use this tool whenever the user asks how to defend against, mitigate,
|
|
97
|
+
or prevent a technique (e.g. "What mitigations exist for T1003?").
|
|
98
|
+
Returns each mitigation's ATT&CK ID (M####), name, and description,
|
|
99
|
+
sourced from the official ATT&CK dataset — do not answer from memory.
|
|
100
|
+
|
|
101
|
+
domain: one of 'enterprise-attack', 'ics-attack', 'mobile-attack'.
|
|
102
|
+
Omit to search all three.
|
|
103
|
+
"""
|
|
104
|
+
return services.get_mitigations_for_technique(attack_id, domain=domain)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def main():
|
|
108
|
+
_preStart()
|
|
109
|
+
logging.info("mitre-attack-mcp server starting..")
|
|
110
|
+
mcp.run()
|
|
111
|
+
logging.info("mitre-attack-mcp server started..")
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
if __name__ == "__main__":
|
|
115
|
+
main()
|
|
@@ -0,0 +1,160 @@
|
|
|
1
|
+
"""Pure business logic. This module must NOT import `mcp`."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
from . import data
|
|
5
|
+
import logging
|
|
6
|
+
import sys
|
|
7
|
+
|
|
8
|
+
logger = logging.getLogger(__name__)
|
|
9
|
+
logging.basicConfig(
|
|
10
|
+
level=logging.INFO,
|
|
11
|
+
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
|
12
|
+
stream=sys.stderr, # never stdout — that's the JSON-RPC channel
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def get_data_stats(*, domain: str | None = None) -> dict | list[dict]:
|
|
17
|
+
"""Counts of techniques, groups, software, and mitigations.
|
|
18
|
+
|
|
19
|
+
domain: restrict to one domain, or omit for a per-domain breakdown.
|
|
20
|
+
"""
|
|
21
|
+
if domain is not None:
|
|
22
|
+
idx = data.get_index(domain)
|
|
23
|
+
return {"domain": domain, **idx.counts()}
|
|
24
|
+
|
|
25
|
+
result = [{"domain": d, **data.get_index(d).counts()} for d in data.DOMAINS]
|
|
26
|
+
logging.info("results %s=", result)
|
|
27
|
+
return result
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def get_technique_by_id(
|
|
31
|
+
attack_id: str,
|
|
32
|
+
*,
|
|
33
|
+
domain: str | None = None,
|
|
34
|
+
) -> dict:
|
|
35
|
+
"""Look up a technique by ATT&CK ID.
|
|
36
|
+
|
|
37
|
+
If `domain` is omitted, all three domains are searched. The returned dict
|
|
38
|
+
includes `domain` so the caller knows where the match came from, and
|
|
39
|
+
`other_domains` lists any additional domains where the same ID also
|
|
40
|
+
exists (ATT&CK IDs are not globally unique across domains).
|
|
41
|
+
"""
|
|
42
|
+
matches: list[tuple[str, dict]] = []
|
|
43
|
+
for d in (domain,) if domain else data.DOMAINS:
|
|
44
|
+
obj = data.get_index(d).get_technique(attack_id)
|
|
45
|
+
if obj is not None:
|
|
46
|
+
matches.append((d, obj))
|
|
47
|
+
|
|
48
|
+
if not matches:
|
|
49
|
+
where = domain or "any domain"
|
|
50
|
+
raise ValueError(f"Technique {attack_id!r} not found in {where}")
|
|
51
|
+
|
|
52
|
+
d, obj = matches[0]
|
|
53
|
+
result = {
|
|
54
|
+
"domain": d,
|
|
55
|
+
"id": data.attack_id_of(obj),
|
|
56
|
+
"name": obj.get("name"),
|
|
57
|
+
"description": (obj.get("description") or "")[:800],
|
|
58
|
+
"tactics": data.tactics_of(obj),
|
|
59
|
+
"platforms": obj.get("x_mitre_platforms", []),
|
|
60
|
+
"is_subtechnique": obj.get("x_mitre_is_subtechnique", False),
|
|
61
|
+
}
|
|
62
|
+
if len(matches) > 1:
|
|
63
|
+
result["other_domains"] = [m[0] for m in matches[1:]]
|
|
64
|
+
result["note"] = (
|
|
65
|
+
"This ID also exists in other domains with different content; "
|
|
66
|
+
"pass `domain` to disambiguate."
|
|
67
|
+
)
|
|
68
|
+
logging.info("results %s=", result)
|
|
69
|
+
return result
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def search_techniques(
|
|
73
|
+
keyword: str,
|
|
74
|
+
*,
|
|
75
|
+
domain: str | None = None,
|
|
76
|
+
limit: int = 10,
|
|
77
|
+
) -> list[dict]:
|
|
78
|
+
"""Search techniques by keyword. `domain=None` searches all domains."""
|
|
79
|
+
domains = (domain,) if domain else data.DOMAINS
|
|
80
|
+
for d in domains:
|
|
81
|
+
if d not in data.DOMAINS:
|
|
82
|
+
raise ValueError(f"Unknown domain {d!r}; expected one of {data.DOMAINS}")
|
|
83
|
+
|
|
84
|
+
results: list[dict] = []
|
|
85
|
+
for d in domains:
|
|
86
|
+
hits = data.get_index(d).search_techniques(keyword, limit=limit)
|
|
87
|
+
results.extend(
|
|
88
|
+
{"domain": d, "id": data.attack_id_of(o), "name": o.get("name")}
|
|
89
|
+
for o in hits
|
|
90
|
+
)
|
|
91
|
+
if len(results) >= limit:
|
|
92
|
+
break
|
|
93
|
+
|
|
94
|
+
logging.info("results %s=", results[:limit])
|
|
95
|
+
return results[:limit]
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def get_group_by_name(
|
|
99
|
+
name: str,
|
|
100
|
+
*,
|
|
101
|
+
domain: str | None = None,
|
|
102
|
+
) -> dict:
|
|
103
|
+
"""Look up a threat group by name or alias.
|
|
104
|
+
|
|
105
|
+
If `domain` is omitted, all three domains are searched; `domain` in the
|
|
106
|
+
result tells the caller where the match was found.
|
|
107
|
+
"""
|
|
108
|
+
matches: list[tuple[str, dict]] = []
|
|
109
|
+
for d in (domain,) if domain else data.DOMAINS:
|
|
110
|
+
obj = data.get_index(d).get_group(name)
|
|
111
|
+
if obj is not None:
|
|
112
|
+
matches.append((d, obj))
|
|
113
|
+
|
|
114
|
+
if not matches:
|
|
115
|
+
where = domain or "any domain"
|
|
116
|
+
raise ValueError(f"Group {name!r} not found in {where}")
|
|
117
|
+
|
|
118
|
+
d, obj = matches[0]
|
|
119
|
+
result = {
|
|
120
|
+
"domain": d,
|
|
121
|
+
"name": obj.get("name"),
|
|
122
|
+
"aliases": obj.get("aliases", []),
|
|
123
|
+
"description": (obj.get("description") or "")[:800],
|
|
124
|
+
}
|
|
125
|
+
if len(matches) > 1:
|
|
126
|
+
result["other_domains"] = [m[0] for m in matches[1:]]
|
|
127
|
+
|
|
128
|
+
logging.info("results %s=", result)
|
|
129
|
+
return result
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
def get_mitigations_for_technique(
|
|
133
|
+
attack_id: str,
|
|
134
|
+
*,
|
|
135
|
+
domain: str | None = None,
|
|
136
|
+
) -> dict:
|
|
137
|
+
"""Return mitigations linked to a technique.
|
|
138
|
+
|
|
139
|
+
If `domain` is omitted, all three domains are searched.
|
|
140
|
+
"""
|
|
141
|
+
matches: list[tuple[str, dict]] = []
|
|
142
|
+
for d in (domain,) if domain else data.DOMAINS:
|
|
143
|
+
obj = data.get_index(d).get_technique(attack_id)
|
|
144
|
+
if obj is not None:
|
|
145
|
+
matches.append((d, obj))
|
|
146
|
+
|
|
147
|
+
if not matches:
|
|
148
|
+
where = domain or "any domain"
|
|
149
|
+
raise ValueError(f"Technique {attack_id!r} not found in {where}")
|
|
150
|
+
|
|
151
|
+
d, tech = matches[0]
|
|
152
|
+
mitigations = data.get_index(d).mitigations_for_technique(tech["id"])
|
|
153
|
+
result = {
|
|
154
|
+
"domain": d,
|
|
155
|
+
"technique_id": data.attack_id_of(tech),
|
|
156
|
+
"technique_name": tech.get("name"),
|
|
157
|
+
"mitigations": mitigations,
|
|
158
|
+
}
|
|
159
|
+
logging.info("results %s=", result)
|
|
160
|
+
return result
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
import pytest
|
|
2
|
+
import json
|
|
3
|
+
from pytest_mcp import assert_tool_call, assert_tool_error
|
|
4
|
+
from mitre_attack_mcp.server import mcp # Import your MCPServer instance
|
|
5
|
+
|
|
6
|
+
# Tell pytest-mcptest which server to test
|
|
7
|
+
pytestmark = pytest.mark.mcp_server(mcp)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def _payload(result):
|
|
11
|
+
"""Extract the tool's return value from a CallToolResult."""
|
|
12
|
+
if getattr(result, "structured_content", None):
|
|
13
|
+
return result.structured_content
|
|
14
|
+
# Fallback: the SDK wraps a bare dict return in {"result": ...}
|
|
15
|
+
text = result.content[0].text
|
|
16
|
+
return json.loads(text)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
async def test_technique_lookup_without_domain(mcp_client):
|
|
20
|
+
result = await assert_tool_call(
|
|
21
|
+
mcp_client, "get_technique_by_id", {"attack_id": "T1059"}
|
|
22
|
+
)
|
|
23
|
+
data = _payload(result)
|
|
24
|
+
if set(data.keys()) == {"result"}:
|
|
25
|
+
data = data["result"]
|
|
26
|
+
assert data["domain"] == "enterprise-attack"
|
|
27
|
+
assert "Command and Scripting" in data["name"]
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
async def test_mitigations_for_technique(mcp_client):
|
|
31
|
+
result = await assert_tool_call(
|
|
32
|
+
mcp_client, "get_mitigations_for_technique", {"attack_id": "T1003"}
|
|
33
|
+
)
|
|
34
|
+
data = _payload(result)
|
|
35
|
+
if set(data.keys()) == {"result"}:
|
|
36
|
+
data = data["result"]
|
|
37
|
+
ids = {m["id"] for m in data["mitigations"]}
|
|
38
|
+
assert {"M1015", "M1043", "M1026"} <= ids
|