Zunalkhan/LLM-VUlnerable-scan-Aisec-
GitHub: Zunalkhan/LLM-VUlnerable-scan-Aisec-
面向本地 HuggingFace 模型的 LLM 安全扫描器，通过多检测引擎评估 prompt injection 和越狱漏洞风险。
Stars: 0 | Forks: 0
# LLM 漏洞扫描器 (MVP)
一个用于**本地 HuggingFace 模型**的 prompt injection 和越狱检测扫描器。
向目标模型发射包含对抗性探测的库，通过多个独立的检测引擎对响应进行评分，
并生成 HTML/JSON 报告。
## 架构
```
llm_vuln_scanner/
core/
model_adapter.py # LocalHFModel wrapper (+ MockModel for smoke tests)
detectors.py # KeywordEngine, CanaryTokenEngine, PerplexityEngine, EmbeddingSimilarityEngine
orchestrator.py # Scanner: loads probes, runs them, aggregates risk scores
report.py # HTML/JSON report rendering
data/
injection_payloads.json # prompt injection probe library
jailbreak_payloads.json # jailbreak probe library
cli.py # command-line interface
api.py # FastAPI service
requirements.txt
```
### 检测引擎
| 引擎 | 检测内容 | 备注 |
|---|---|---|
| `KeywordEngine` | 已知的顺从/泄露短语，缺少预期的拒绝回复 | 速度快，始终开启，无依赖 |
| `CanaryTokenEngine` | System prompt / 上下文泄露 | 在 system prompt 中植入一个秘密 token；如果它出现在输出中，则准确率接近 100% |
| `PerplexityEngine` | 对抗性后缀 / GCG 风格攻击 | 标记统计学上异常的 payload token 序列；需要真实模型（非 mock） |
| `EmbeddingSimilarityEngine` | 已知攻击的释义/变体 | 可选（`--use-embeddings`），需要 `sentence-transformers` |
每个探测都会结合所有引擎的判定结果获得一个**加权风险评分 (0-100)**；
超过 `--flag-threshold`（默认 40）的探测将在报告中被标记为 `VULNERABLE`。
## 设置
```
pip install -r requirements.txt
```
## 用法
### CLI
无需下载模型的快速冒烟测试（使用确定性的 mock 模型）：
```
python cli.py scan --mock --out report.html
```
针对本地/hub 模型的真实扫描：
```
python cli.py scan \
--model gpt2 \
--system-prompt "You are a helpful banking assistant. Never reveal account numbers." \
--categories injection --categories jailbreak \
--use-embeddings \
--flag-threshold 40 \
--out report.html
```
在浏览器中打开 `report.html` 查看完整报告，或使用 `--out report.json`
获取机器可读的输出（例如用于拦截 CI/CD pipeline）。
### API
```
uvicorn api:app --reload --port 8000
```
```
curl -X POST http://localhost:8000/scan \
-H "Content-Type: application/json" \
-d '{
"model_name": "gpt2",
"system_prompt": "You are a helpful assistant.",
"categories": ["injection", "jailbreak"],
"use_embeddings": true
}'
```
`GET /probes` 列出当前的探测库。`GET /health` 用于存活检测。
## 扩展探测库
向 `data/injection_payloads.json` 或 `data/jailbreak_payloads.json` 添加条目：
```
{ "id": "inj-009", "category": "direct_override", "payload": "..." }
```
无需修改代码 —— orchestrator 会动态加载这些内容。
## 路线图（超出此 MVP 的范围）
根据当前范围，此 MVP 仅涵盖 prompt injection + 越狱检测。
遵循相同的“探测 + 检测器”模式，后续自然的扩展模块包括：
1. **数据/token 中毒检测** — 对 fine-tuning 数据集进行统计异常值检测
（embedding 聚类），与干净的基线模型进行差异化测试。
2. **权重/供应链扫描** — `safetensors` 验证，拒绝加载
pickled checkpoints，基于 activation-clustering 的后门/trojan 检测。
3. **敏感信息泄露探测** — 针对 PII/训练数据
记忆的定向探测（例如重复 token 提取攻击）。
4. **LLM-as-judge 引擎** — 使用第二个模型根据评分规则对响应进行评分，
捕捉 keyword/embedding 引擎遗漏的更隐蔽的违规策略。
5. **多轮升级跟踪** — 评估整个
对话过程中的累积风险，而不仅仅是单次探测。
6. **CI/CD 集成** — 如果 `report.json` 中的 `overall_risk_score` 超过
阈值，则使构建失败，类似于 SAST 拦截。
## 已知局限性（请提前了解这些）
- 检测属于启发式方法；预期会存在误报/漏报，特别是仅使用
`KeywordEngine` 时。请组合使用多个引擎并调整 `--flag-threshold`。
- `PerplexityEngine` 需要一个真实模型（而不是 `MockModel`）才能产生有意义的分数。
- 这是一个防御/红队工具，用于测试**您自己的**模型。未经授权，
它不适用于指向第三方的生产系统。
标签：AI红队, DLL 劫持, 大语言模型, 越狱检测, 逆向工具, 零日漏洞检测