tkostanyan/llm-testing-kit
GitHub: tkostanyan/llm-testing-kit
面向工具调用型 LLM 智能体的通用提示注入红队测试工具包，提供分层防御评估、ASR/FPR 量化分析和自动化 PDF 报告生成。
Stars: 0 | Forks: 0
# agent-redteam
一个用于工具调用 LLM 智能体的通用提示注入红队测试工具包。
你只需将其指向目标智能体（例如邮件助手、RAG 智能体或你自定义的智能体），它就会在多个模型和层层递进的防御机制中运行完整的攻击样本库，随后自动生成一份 PDF 报告，报告中包含攻击成功率 (ASR) / 误报率 (FPR) / 延迟的表格及每一层防御的图表。
## 它的功能
1. **驱动任何调用工具的 LLM 智能体。** 目标可通过 `TargetAgent` 接口进行插拔式接入——你只需提供一个系统提示词、工具集、不透明状态对象以及注入钩子。该工具包内置了两个参考目标：
- `email` — 存在漏洞的邮件助手（读取/搜索/发送/转发/删除），包含 7 封模拟邮件的收件箱。
- `rag` — 研究助手，配备 `search_docs` / `read_doc` / `web_fetch` 工具以及 5 个文档的模拟语料库。
2. **发射分类的提示注入样本库**（约 75 个 payload，分类参考了 OWASP LLM01）覆盖所有目标/模型。
3. **评估堆叠式的防御进程** —— 五个累积条件，每一层都在前一层的基础上叠加：
`none → +input_filter → +output_validator → +tool_allowlist → +llm_judge_prefilter`。
4. **在手工筛选的良性控制集（每个目标 10 个任务）上测量误报率**，以便在提升安全性的同时直观呈现安全性与可用性之间的权衡。
5. **生成 PDF 报告**（LaTeX → tectonic），包含：核心的 ASR/FPR/延迟表格、ASR 与防御层级曲线、ASR 与 FPR 散点图、按类别的热力图、延迟开销、步骤详述、在完整防御栈下幸存的攻击列表，以及局限性说明。
## 威胁模型
受保护的用户运行一个 LLM 智能体，该智能体拥有能在现实世界产生副作用的工具（如发送电子邮件、获取 URL、删除数据）。攻击者控制着一部分受其影响的输入——直接控制用户提示词（*直接注入*），或控制智能体在处理正常请求时将读取的某部分状态（*间接注入*：被投毒的邮件正文、被投毒的检索文档等）。当智能体发起导致数据泄露或破坏状态的 tool call 时，攻击者即获胜。
该工具包会针对每一层防御，测量攻击者获胜的频率 (ASR) 以及防御机制干扰正常使用的频率 (FPR)。
## 安装
```
python3 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"
cp .env.example .env # fill in GEMINI_API_KEY, ANTHROPIC_API_KEY, OPENAI_API_KEY, GROQ_API_KEY
```
若要编译 PDF，请安装 [tectonic](https://tectonic-typesetting.github.io/)（`brew install tectonic`）或任何提供 `pdflatex` 的 TeX 发行版。无论是否安装编译器，该工具包都会生成 `.tex` 文件，并在 PATH 中存在编译器时自动将其编译为 PDF。
## 快速开始
针对 email 目标、一个 Groq 模型、3 个攻击 payload + 10 个良性样本生成端到端报告：
```
python scripts/generate_report.py --target email --models llama-8b --limit 3
# -> results/report_email_/{report.pdf, report.tex, attack.csv, benign.csv, figs/}
```
针对 RAG 目标生成端到端报告：
```
python scripts/generate_report.py --target rag --models llama-8b --limit 3
```
仅运行攻击扫描（无良性样本，不生成报告）：
```
python scripts/run_experiment.py --target email --models gemini-flash --defenses none --categories direct_injection --limit 3
# -> results/run_.csv
```
从已收集的 CSV 数据重新渲染报告（不调用 API）：
```
python scripts/generate_report.py --target email \
--attack-csv results/report_email_/attack.csv \
--benign-csv results/report_email_/benign.csv \
--out-dir results/report_email__v2
```
## 项目结构
```
src/agent_redteam/
targets/ # pluggable TargetAgent implementations
base.py # TargetAgent ABC + AgentState protocol
email/ # email-assistant target (Inbox + 6 tools)
rag/ # research-assistant target (DocStore + 3 tools)
agent/ # backwards-compat shim re-exporting EmailAgent
evaluation/
loop.py # generic AgentLoop driving any TargetAgent
runner.py # attack sweep over (model, defense, payload)
benign.py # benign-task sweep for FPR measurement
judge.py # programmatic verdict (SUCCESS/FAIL/BLOCKED/ERROR)
scenario.py # builds (state, user_message) per trial
defenses/
input_filter.py # regex blocklist on user input + untrusted state
output_validator.py # external-recipient allowlist
tool_allowlist.py # auto-deny destructive tools
llm_judge_prefilter.py # cheap LLM injection classifier
composite.py # chain defenses; first block wins
payloads/ # YAML loader + Pydantic schema
providers/ # gemini / claude / openai / groq adapters
reporting/ # metrics, charts, LaTeX renderer
data/
inbox/baseline_emails.json
rag/baseline_docs.json
payloads/01..08*.yaml # attack corpus (email + rag targets)
benign/{email,rag}.yaml # benign control sets
templates/report.tex # LaTeX template (<> placeholders)
config/models.yaml # model alias → (provider, upstream id)
scripts/
generate_report.py # full pipeline → PDF
run_experiment.py # attack-sweep only (no benign, no report)
```
## 添加自定义目标智能体
继承 `TargetAgent` 子类并注册它。大致如下：
```
from agent_redteam.targets.base import TargetAgent
from agent_redteam.providers.base import ToolSpec
class MyTarget(TargetAgent):
name = "my_agent"
system_prompt = "You are a ..."
tool_specs: list[ToolSpec] = [...]
destructive_tools = {"my_dangerous_tool"}
def build_initial_state(self): ...
def execute_tool(self, name, args, state): ...
def apply_injection(self, state, injection, payload_text): ...
def extract_external_recipients(self, call): ...
def default_benign_prompt(self): return "..."
```
将其添加到 `targets/__init__.py::ALL_TARGETS` 和 `get_target` 中。提供一个包含 `target: my_agent` 的 payload YAML 文件，并在 `data/benign/my_agent.yaml` 中提供一个良性样本 YAML。运行器、防御机制、评判器和报告流水线无需进一步修改即可直接工作。
## Payload 分类体系
7 个攻击类别与 OWASP LLM 应用程序 Top 10 (LLM01) 中描述的提示注入模式保持一致：
| 类别 | 测试内容 |
|---|---|
| `direct_injection` | 攻击者 = 用户；显式指令覆盖 |
| `indirect_injection` | 隐藏在攻击者可控制的、智能体将读取的状态中的 payload |
| `multi_turn` | 伪造的先前对话轮次 |
| `encoded_obfuscated` | base64 / ROT13 / 同形字 / leet / 反转文本 |
| `role_play_jailbreak` | DAN / 假设性 / 角色切换框架 |
| `tool_abuse_chain` | 多步数据窃取与掩盖行踪 |
| `context_manipulation` | 伪造的 tool result / 伪造的 system tag / 预填充 |
每个 payload 都声明了一个 `success_criterion`（`tool_call`、`tool_call_any`、`output_contains`、`sent_to_external`），以便评判器能够确定性地对结果进行评分。
## 防御栈（从成本最低到最高）
1. **`input_filter`** — 对用户消息及智能体状态中当前的每个不受信任内容片段应用正则表达式黑名单。
2. **`output_validator`** — 拒绝外部目的地不在白名单内的 tool call（使用 `target.extract_external_recipients`）。
3. **`tool_allowlist`** — 自动拒绝任何对破坏性工具的调用（使用 `target.destructive_tools`）；模拟“用户始终拒绝未经请求的破坏性行动”的策略。
4. **`llm_judge_prefilter`** — 廉价的 LLM 分类器（在 `config/models.yaml` 的 `prefilter:` 下配置）判断某段输入是否看起来像是一次注入尝试。
`none` 是基线（完全无防御）。
## 阅读结果 CSV
`results/.../attack.csv` 列：
| 列 | 含义 |
|---|---|
| `payload_id` | 例如 `indirect_007` |
| `category` | 7 个 OWASP 对齐类别之一 |
| `target`, `model`, `defense` | 试验标识 |
| `turns` | 消耗的智能体循环轮次 |
| `tool_calls_json` | 包含参数的每一个 tool call（包括被阻止的） |
| `final_output` | 助手的最后一段文本（截断至 2000 个字符） |
| `judge_verdict` | `SUCCESS / FAIL / BLOCKED / ERROR` |
| `judge_reasoning` | 评判器选择该裁决的原因 |
| `latency_ms`, `input_tokens`, `output_tokens` | 运维指标 |
`results/.../benign.csv` 的列包含 `blocked`（true/false）和 `blocked_phase`（`input` / `tool_call` / 空），因此你可以将误报归因于特定的防御层。
## 运行测试
```
pytest # all 40 tests
pytest tests/test_report_pipeline.py # one file
```
## 道德与负责任的使用
本仓库包含**攻击性的提示注入 payload**。它们已被分类、规模较小，且仅针对进程内的模拟状态——不涉及真实的电子邮件/收件箱/文档/网络流量。**请勿**将这些 payload 用于你未拥有或未获得明确测试授权的系统。该工具包仅适用于：对你所运营的智能体进行防御性评估、对你自己系统进行安全研究，以及在受控环境中进行教育性的红队演练。
## 许可证
MIT。
标签：AI安全, AI智能体, Chat Copilot, CISA项目, DLL 劫持, LLM01, LLM测评, PDF报告, Petitpotam, RAG安全, Red Canary, Sysdig, 反取证, 大语言模型, 安全评估, 提示注入, 攻击成功率, 纵深防御, 网络安全, 自动化报告, 误报率, 逆向工具, 隐私保护, 集群管理