bellakhalrania/AI-Threat-Intelligence---Red-Team
GitHub: bellakhalrania/AI-Threat-Intelligence---Red-Team
基于 MITRE ATT&CK 真实数据的 RAG 平台,通过多个专职 AI Agent 为安全分析人员提供防幻觉的威胁情报查询与归因分析服务。
Stars: 0 | Forks: 0
# MITRE ATT&CK RAG / 红队情报平台
一个检索增强生成平台,允许 SOC 分析师、威胁
猎手和红队操作员通过四个专门的 AI agent 查询 MITRE ATT&CK 知识库,
该平台基于真实的 ATT&CK 数据(而非 LLM
记忆)构建,并由 Groq 托管的 Llama 3.3 提供服务。
此实现已基于**真实的、实时的
MITRE ATT&CK Enterprise STIX bundle**(包含 25,842 个对象,直接从
`github.com/mitre/cti` 获取)构建并测试——本仓库中的每个解析器函数都已在
实际生产数据上运行,而非合成样本。
## 1. 架构
```
┌─────────────────────────────────────────────────────────────────────┐
│ Layer 1 — Data Source │
│ MITRE CTI repo (github.com/mitre/cti) → enterprise-attack.json │
│ STIX 2.1 bundle: 858 attack-patterns, 21k relationships, │
│ 729 malware, 189 intrusion-sets, 95 tools, 268 mitigations, │
│ 1758 analytics, 699 detection strategies │
└───────────────────────────────┬────────────────────────────────────-┘
▼
┌─────────────────────────────────────────────────────────────────────┐
│ Layer 2 — Data Processing (parser/stix_parser.py) │
│ Loads bundle → builds id-lookup table → filters revoked/deprecated │
│ → emits 4 datasets: │
│ • techniques (attack-pattern) → attack_db │
│ • procedures (relationship type="uses") → redteam_db │
│ • actors_tools (intrusion-set/malware/tool) → actor_db │
│ • detections (course-of-action/analytic/ → detection_db │
│ detection-strategy) │
└───────────────────────────────┬────────────────────────────────────-┘
▼
┌─────────────────────────────────────────────────────────────────────┐
│ Layer 3 — Embedding Pipeline (embeddings/embedding_pipeline.py) │
│ BAAI/bge-base-en-v1.5 (768-dim, retrieval-optimized) │
│ Documents embedded as-is; queries get the bge retrieval prefix │
└───────────────────────────────┬────────────────────────────────────-┘
▼
┌─────────────────────────────────────────────────────────────────────┐
│ Layer 4 — Vector DB (vectordb/chroma_manager.py) │
│ ChromaDB PersistentClient, 4 cosine-similarity collections, │
│ local disk storage at storage/chroma_db/, no external DB needed │
└───────────────────────────────┬────────────────────────────────────-┘
▼
┌─────────────────────────────────────────────────────────────────────┐
│ Layer 5 — Agent Layer (agents/, orchestrator.py) │
│ Red Team Agent → attack_db + redteam_db │
│ CTI Analyst Agent → attack_db + detection_db (healthcare framing) │
│ Attribution Agent → actor_db (confidence-scored, never invents) │
│ Detection Agent → detection_db + attack_db │
│ Orchestrator routes free-text queries to the right agent │
└─────────────────────────────────────────────────────────────────────┘
```
**每次查询的 RAG 工作流:** 用户查询 → ID 验证防护机制 →
嵌入查询(bge prefix)→ 跨 agent 的集合进行向量搜索 →
按余弦距离合并和排序 → 构建上下文块 →
带有基础规则系统提示词的 Groq Llama 3.3 → 将答案和引用来源返回给调用者。
**防止幻觉(`validation/id_validator.py`):** 用户查询中
每个符合 MITRE-ID 格式的 token(正则表达式:`T1003`, `T1003.001`, `TA0001`,
`S0154`, `G0016`, `M1015`)在查询到达
LLM 之前,都会在 ChromaDB 中进行**精确元数据匹配**检查。如果摄取的
数据集中不存在某个 ID,agent 将返回拒绝响应——LLM 永远不会
有机会“解释”一个虚假的技术。这已在真实的
数据集上进行了测试:`T1047`(真实)→ 有效,`T9999`(虚假)→ 正确被拒绝。
## 2. 项目布局
```
mitre_rag_platform/
├── config.py # all paths, model names, tunables
├── requirements.txt
├── ingest.py # one-shot: download → parse → embed → store
├── app.py # CLI chat interface
├── web_app.py # Streamlit web interface
├── orchestrator.py # query routing across the 4 agents
├── data/
│ ├── download_mitre.py # pulls enterprise-attack.json from GitHub
│ └── parsed/ # generated: 4 JSON datasets (after ingest)
├── parser/
│ └── stix_parser.py # STIX → 4 structured datasets
├── embeddings/
│ └── embedding_pipeline.py # BAAI/bge-base-en-v1.5 wrapper
├── vectordb/
│ └── chroma_manager.py # ChromaDB collections, ingest + query
├── validation/
│ └── id_validator.py # hallucination-prevention guardrail
├── llm/
│ └── groq_client.py # Groq API call + grounding system prompt
├── agents/
│ ├── base_agent.py # shared RAG pipeline (retrieve + answer)
│ ├── redteam_agent.py
│ ├── cti_agent.py
│ ├── attribution_agent.py # adds confidence scoring
│ └── detection_agent.py
└── storage/chroma_db/ # persistent vector DB (created on first run)
```
## 3. 设置
```
cd mitre_rag_platform
python3 -m venv venv && source venv/bin/activate
pip install -r requirements.txt
export GROQ_API_KEY=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 # https://console.groq.com/keys
```
### 构建知识库(一次性操作,在 CPU 上约需 5-10 分钟)
```
python ingest.py
```
这将下载实时的 MITRE bundle,将其解析为 4 个数据集,
使用 `BAAI/bge-base-en-v1.5` 嵌入每个文档,并将所有内容
存储在 `storage/chroma_db/` 中。每当 MITRE 发布新的 ATT&CK
版本(每年几次)时重新运行此操作,以保持知识库最新。
为了实现快速的开发循环,请先摄取一小部分数据:
```
python ingest.py --limit 50
```
### 运行 CLI 聊天界面
```
python app.py
```
```
============================================================
MITRE ATT&CK RAG / Red Team Intelligence Platform
Agents: [redteam] [cti] [attribution] [detection] [auto]
============================================================
(auto)> How does Ryuk gain persistence?
[Red Team Agent]
...grounded answer citing specific T-IDs from retrieved context...
(auto)> /agent attribution
(attribution)> Which actors use Mimikatz and PsExec together?
[Attribution Agent]
1. – confidence ...
```
### 运行 Web 界面(推荐)
```
streamlit run web_app.py
```
在浏览器中打开 `http://localhost:8501`
**功能:**
- **多 agent 聊天** —— 手动选择 agent 或使用自动路由
- **持久聊天历史** —— 所有查询和响应都存储在会话中
- **检索来源显示** —— 查看所有检索到的分块,包括:
- 余弦距离分数
- 置信度级别(用于归因 agent)
- 集合元数据
- **数据库统计** —— 实时查看已摄取的集合
- **Agent 选择侧边栏** —— 固定到特定 agent 或启用自动路由
- **响应式设计** —— 适用于桌面设备和平板电脑
**示例查询:**
```
"How does Ryuk gain persistence?" → Red Team Agent (auto-routed)
"Detect Mimikatz usage in logs" → Detection Agent
"Which APT groups target healthcare?" → CTI Analyst Agent
"Identify actors using T1003 and T1055" → Attribution Agent
```
## 4. 关于此沙盒与您机器的说明
此仓库中的所有内容均在此开发沙盒内基于**真实的 MITRE
数据**构建并验证:
- `parser/stix_parser.py` 已在实际的 46 MB `enterprise-attack.json` 上运行
(提取了 697 项技术、16,903 项过程、995 个参与者/工具、
2,499 条检测记录,并对样本记录进行了质量抽查)。
- `vectordb/chroma_manager.py` 已完成端到端测试:摄取、精确 ID
元数据查找和余弦向量搜索均确认正常工作。
- `validation/id_validator.py` 已在真实数据集上进行测试:
真实 ID(`T1047`, `S0140`, `G0016`, `M1021`)验证为 `True`;虚假 ID
(`T9999`)正确验证为 `False` 并在
进行任何 LLM 调用*之前*产生拒绝消息。
标签:Cloudflare, DLL 劫持, Kubernetes, MITRE ATT&CK, SOC分析, 人工智能, 大语言模型, 威胁情报, 开发者工具, 数据展示, 检索增强生成, 用户模式Hook绕过, 红队, 逆向工具