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绕过, 红队, 逆向工具