spoorthi2615/HuntGPT
GitHub: spoorthi2615/HuntGPT
该项目利用微调大语言模型将网络日志自主映射到 MITRE ATT&CK 技术进行威胁狩猎，并通过专门的过滤层防御嵌入在日志中的对抗性提示注入攻击。
Stars: 0 | Forks: 0
# HuntGPT-Shield
**自主网络威胁狩猎与对抗性提示注入防御**
HuntGPT-Shield 是一个全面的、端到端的网络安全 pipeline，旨在使用大型语言模型 (LLM) 将原始网络日志自主映射到 MITRE ATT&CK 技术。至关重要的是，它包含了一个健壮的过滤层，以保护 LLM 免受隐藏在网络 payload 中的对抗性提示注入攻击。
## 🏗️ 架构
该系统由四个主要层组成：
1. **PromptGuard（过滤器）**：一个微调过的 `DeBERTa-v3-small` 序列分类器。它在传入的 Zeek 网络日志到达 LLM 之前，扫描其中是否包含嵌入式提示注入，并返回一个布尔的阻止/允许决定。
2. **ThreatHunter（分析师）**：一个使用 QLoRA 微调的 `Mistral-7B-Instruct-v0.2` 模型。它接收干净的网络日志并输出可解析的 JSON，将活动映射到特定的 MITRE ATT&CK `technique_id` 并生成分析师假设。
3. **自一致性评分器**：通过在 `temperature=0.7` 下采样 N=5 次预测，并返回多数投票和一致率（置信度/ECE），来评估 LLM 的置信度。
4. **React 仪表盘**：一个 Vite/Tailwind 前端，提供 SOC 分析师 UI，用于提交日志批次、查看实时 pipeline 执行情况、探索检测到的技术的 ATT&CK 热力图，以及监控系统的评估指标。
## 🚀 快速开始（演示模式）
你可以在“Zero-Shot Stub Mode”下运行完整的 UI 和 pipeline，而无需 GPU 或微调模型。
### 前置条件
* Docker & Docker Compose
* 运行 `mistral:instruct` 的本地 [Ollama](https://ollama.ai/)（用于 zero-shot LLM 推理）
```
# Clone 仓库
git clone https://github.com/spoorthi2615/HuntGPT.git
cd HuntGPT
# 确保本地正在运行 Ollama
ollama run mistral:instruct
# 启动 frontend 和 backend 服务
docker compose up --build
```
导航到 `http://localhost:5173` 即可访问 SOC Dashboard。
## 🧠 训练模型
为了全面训练 PromptGuard 和 ThreatHunter 模型，您必须按照以下特定顺序执行数据准备脚本和训练任务。
### 1. 数据准备
该 pipeline 依赖于 MITRE ATT&CK STIX 数据库和 DARPA OPTC 数据集（由于体积较大未包含在内）。
```
# 下载并解析 MITRE STIX 数据
python src/data/attack_scraper.py
# 生成 prompt injection 训练语料库（stratified split）
python src/data/injection_builder.py
# 将解析后的 OPTC Zeek 日志与 ground-truth technique labels 进行合并
python src/data/build_training_pairs.py
```
### 2. PromptGuard 训练 (DeBERTa)
可以在本地 GPU 或 Colab T4 上运行（约 20 分钟）。
```
# 训练模型
accelerate launch src/models/train_guard.py
# 评估并生成 Table 1 指标
python src/eval/eval_guard.py
```
### 3. ThreatHunter 微调 (Mistral-7B QLoRA)
需要 Colab A100 或同等配置（约 2 小时）。
1. 将 `data/processed/training_pairs.json` 上传到 Google Drive。
2. 从上到下打开并运行 `notebooks/train_llm_colab.ipynb`。
3. 将生成的 adapter 权重下载到 `models/threathunter/` 中。
### 4. 系统评估与激活
一旦 ThreatHunter adapter 被放入 `models/threathunter/`，系统就会自动检测到它。您可以通过设置 `USE_FINETUNED_MODEL=true` 来强制 pipeline 使用该微调模型。
```
# 运行最终评估以生成 Table 2 和 ECE calibration plot
USE_FINETUNED_MODEL=true python src/eval/eval_threathunter.py
python src/eval/eval_calibration.py
```
## 🧪 运行测试
该项目使用 `pytest` 在整个 pipeline 中强制执行严格的 API 契约。
```
# 安装测试依赖
pip install -r requirements-test.txt
# 运行测试套件
pytest tests/
```
专为微调 adapter 设计的测试将自动 `skip`，直到 `models/` 目录中存在 adapter 产物，从而保持 CI 为绿色通过状态。
标签：AI风险缓解, ATT&CK框架, DLL 劫持, React, Syscalls, 大语言模型, 安全规则引擎, 提示注入防御, 源代码安全, 网络安全, 请求拦截, 逆向工具, 隐私保护