MuhajirMakki/ThreatGraphX
GitHub: MuhajirMakki/ThreatGraphX
将 CVE、CWE、CAPEC 和 ATT&CK 四大安全知识源统一为异构知识图谱，利用图神经网络进行零样本威胁关系预测与可解释攻击路径生成的网络威胁情报系统。
Stars: 0 | Forks: 0
# ThreatGraphX 🕸️🔐
**使用知识图谱和图神经网络的可解释网络威胁情报**
ThreatGraphX 将四大主要网络安全知识源 —— **NVD CVE、MITRE CWE、CAPEC 和 ATT&CK** —— 整合到一个统一的异构知识图谱中，然后使用图神经网络来预测缺失的威胁关系，并解释*为什么*某个漏洞会映射到特定的攻击技术和缓解措施。
ThreatGraphX 无需在四个互不关联的数据库中手动追踪某个漏洞，而是自动遍历整个链路：
```
CVE → CWE → CAPEC → ATT&CK Technique → Mitigation
(vulnerability) (weakness) (attack pattern) (adversary behavior) (defense)
```
## ✨ 核心特性
- **统一知识图谱** —— 包含 7,502 个节点、15,573 条边、5 种关系类型，基于原始的 CVE/CWE/CAPEC/ATT&CK 数据源构建
- **基于 GNN 的链接预测** —— GraphSAGE、R-GCN 以及自定义的 Transformer 增强模型（**ThreatGraphX-Tx**）
- **Zero-Shot 评估** —— 测试模型在预测训练期间从未见过的 CVE 关系时的表现，模拟真实世界中新披露的漏洞
- **两阶段（TxGNN 风格）训练** —— 首先在全图上进行预训练，然后专门针对与 CVE 相关的边进行微调，以获得更强的 Zero-Shot 泛化能力
- **可解释输出** —— 每个预测都附带一条从漏洞到缓解措施的可追溯多跳路径，而不仅仅是一个黑盒分数
## 📊 结果摘要
**标准评估（Transductive）**
| Model | Accuracy | F1 | AUC-ROC |
|---|---|---|---|
| GraphSAGE | 90.65% | 0.9129 | 0.9527 |
| R-GCN | **92.73%** | 0.9306 | **0.9735** |
| ThreatGraphX-Tx | 92.73% | **0.9316** | 0.9716 |
**严格 Zero-Shot 评估（未见过的 CVE 节点）**
| Model | Accuracy | F1 | AUC-ROC |
|---|---|---|---|
| GraphSAGE | 76.06% | 0.7149 | 0.8781 |
| R-GCN | 78.23% | 0.7498 | 0.8995 |
| **ThreatGraphX-Tx (两阶段)** | **92.02%** | **0.9208** | **0.9739** |
为了进行比较，还测试了 Translational 基线模型（TransE、DistMult），但它们的表现非常差（Hits@10 ≈ 0），这证实了这种异构图必须使用具备关系感知能力的 GNN。
## 🏗️ Pipeline 概述
本项目构建为一个 20 步的 Pipeline：
1. **数据检查** —— 验证原始的 CVE/CWE/CAPEC/ATT&CK 文件
2. **实体提取** —— 解析 CVE、CWE、CAPEC、ATT&CK 技术与缓解措施、软件/产品
3. **边构建** —— 构建 has_weakness、mapped_to、maps_to、mitigated_by、affects 关系
4. **知识图谱构建** —— 通过 node2id/relation2id 映射合并为单个图谱
5. **图谱可视化** —— 检查关系和节点类型的分布
6. **数据集划分** —— 70/15/15 的训练/验证/测试集划分，并转换为 (head, relation, tail) 三元组
7. **基线 Embedding** —— TransE 和 DistMult (PyKEEN)
8. **GNN 模型** —— GraphSAGE、R-GCN 和 ThreatGraphX-Tx（Transformer 增强）
9. **Zero-Shot 与严格 Zero-Shot 评估** —— 包含两阶段微调
10. **可解释性模块** —— 为任意 CVE 生成多跳威胁路径
## 📁 使用的数据集
| Source | Format | Role |
|---|---|---|
| [NVD CVE](https://nvd.nist.gov/vuln/data-feeds) | JSON (NVD 2.0) | 漏洞节点 |
| [MITRE CWE](https://cwe.mitre.org/data/downloads.html) | XML (CWE v7.3) | 弱点类别节点 |
| [MITRE CAPEC](https://capec.mitre.org/data/downloads.html) | XML (CAPEC v3.5) | 攻击模式节点 |
| [MITRE ATT&CK](https://attack.mitre.org) | STIX 2.1 JSON | 技术与缓解节点 |
## 🧠 模型
| Model | Type | Highlight |
|---|---|---|
| TransE / DistMult | Translational KG embedding | 基线模型 —— 对于这种异构图而言表现不足 |
| GraphSAGE | Neighborhood aggregation GNN | 召回率高，精确率较低 |
| R-GCN | Relation-aware GNN | 标准评估中表现最好 |
| **ThreatGraphX-Tx** | R-GCN embeddings + Transformer encoder | 对未见过的 CVE 具有最佳的 Zero-Shot 泛化能力 |
## 🔍 可解释性示例
对于 `CVE-2026-12212`，ThreatGraphX 追踪到以下威胁路径：
```
CVE-2026-12212 → CWE-266 (Incorrect Privilege Assignment)
→ CAPEC-19 (Embedding Scripts within Scripts)
→ T1037 (Boot or Logon Initialization Scripts)
→ Mitigation: Execution Prevention
```
这让分析师能够直观地看到**为什么**某项技术是相关的 —— 而不仅仅是一个置信度分数。
## 🛠️ 技术栈
- Python、PyTorch、PyTorch Geometric
- PyKEEN（TransE、DistMult 基线模型）
- NetworkX（图谱构建与可视化）
- Pandas / NumPy / scikit-learn
## 📂 仓库结构
```
data/
├── raw/ # Original CVE, CWE, CAPEC, ATT&CK source files
├── processed/ # Extracted node & edge CSVs
└── splits/ # Train/val/test edge & triple files
notebooks/ # Step-by-step pipeline notebooks (Steps 01–20)
results/ # Model metrics, threat discovery tables
```
## 📚 参考文献
1. Usama, M., Nagaty, K. and Elmaghawry, N., 2025. *Construction of a Unified Knowledge Graph for Cyber Threat Intelligence.* ICCA 2025.
2. Liang, K., Liu, Y., Zhou, S., 2024. *Knowledge graph contrastive learning based on relation-symmetrical structure.* IEEE TKDE.
3. Ji, S., Pan, S., Cambria, E., Marttinen, P., Yu, P.S., 2021. *A survey on knowledge graphs.* IEEE TNNLS.
4. Alsaheel, S. et al., 2021. *ATT&CK-KG: A knowledge graph of adversary techniques and tactics.* IEEE Big Data.
## 👤 作者
**Abdu Samad Shah**
提交给 Dr. Muhammad Ishaq（助理教授），IM Sciences
标签：凭据扫描, 可解释AI, 图神经网络, 威胁情报, 安全, 开发者工具, 漏洞分析, 特权检测, 超时处理, 路径探测, 逆向工具