urcuqui/AgentBreak
GitHub: urcuqui/AgentBreak
一个演示 AI Agent 间接 prompt 注入攻击与确定性防御措施的教学实验项目。
Stars: 0 | Forks: 0
# 使用 Python 黑客攻击 AI Agent
这是一个 7-8 分钟的会议演示，展示了存在漏洞的 AI Agent 如何通过 RAG pipeline 检索到的文档内的**间接 prompt 注入**被操纵，以及如何通过确定性的 Python 控制来阻止*相同*的攻击。
## 教学目标
在现场观众面前具体说明：
1. 间接 prompt 注入是 Agent 系统的一种真实故障模式。
2. 将检索到的内容视为受信任的指令是根本原因。
3. 缓解措施不是“客气地请求模型”——而是确定性的 Python：工具 allowlist、参数验证、egress 控制、将指令与不受信任的上下文分离，以及结构化的审计日志。
## 架构
```
+---------------------------+
user ---> | UserRequest |
| question + task |
+-------------+-------------+
|
v
+---------------------------+
| RAG (local, keyword) |
| trust_level=untrusted |
+-------------+-------------+
|
v
+---------------------------+
| Simulated agent brain | (deterministic LLM stand-in)
+-------------+-------------+
|
proposed ToolCall(s)
|
VULNERABLE PATH | SECURE PATH
no checks v policy layer
+---------------------------+
| authorize_tool |
| validate_arguments |
| check_egress |
+-------------+-------------+
|
v
+---------------------------+
| Tool runtime (local only) |
+---------------------------+
```
## 项目布局
```
AgentBreak/
├── README.md
├── pyproject.toml
├── requirements.txt
├── data/
│ ├── clean_knowledge_base.json
│ ├── poisoned_knowledge_base.json
│ └── customers.json
├── src/hacking_ai_agents/
│ ├── cli.py
│ ├── models.py
│ ├── rag.py
│ ├── tools.py
│ ├── agents_common.py
│ ├── vulnerable_agent.py
│ ├── secure_agent.py
│ ├── policies.py
│ ├── security.py
│ ├── presenter.py
│ ├── llm.py ← Ollama client + deterministic simulator
│ ├── webapp.py ← FastAPI app (SSE streaming)
│ ├── web/ ← single-page UI (HTML/CSS/JS)
│ └── logging_utils.py
├── tests/
└── scripts/
├── run_demo.sh
└── run_demo.ps1
```
## 安装
需要 Python 3.11+。运行时不需要网络访问。
```
python -m venv .venv
source .venv/bin/activate # Windows: .venv\Scripts\Activate.ps1
pip install -e ".[dev]"
```
## 可选：通过 Ollama 使用本地 LLM
该演示开箱即用，带有一个确定性的**模拟器**，因此它永远不会在台上卡住。要使用真实的本地模型驱动 agent，请安装 [Ollama](https://ollama.com) 并拉取一个小型 instruct 模型：
```
# 1. 安装 Ollama（参见 https://ollama.com）
# 2. 启动 server（安装后通常作为 service 运行）
ollama serve &
# 3. 拉取一个小型 model
ollama pull llama3.2:3b
# 4. 验证是否已安装
ollama list
```
如果出现以下情况，CLI 和 Web 应用将自动回退到模拟器：
* 无法在 `http://localhost:11434` 访问 Ollama，或者
* 尚未拉取请求的模型标签，或者
* 模型返回不可解析的 JSON。
使用环境变量或标志覆盖默认值：
```
export OLLAMA_BASE_URL=http://localhost:11434
export OLLAMA_MODEL=llama3.2:3b
```
## CLI 命令
```
python -m hacking_ai_agents.cli normal
python -m hacking_ai_agents.cli attack
python -m hacking_ai_agents.cli secure
python -m hacking_ai_agents.cli full-demo
python -m hacking_ai_agents.cli full-demo --no-pause # automatic
python -m hacking_ai_agents.cli attack --use-llm # try the local LLM
python -m hacking_ai_agents.cli attack --use-llm --model llama3.2:3b
```
常用标志：`--verbose`、`--no-color`、`--no-pause`（仅在 `full-demo` 上可用）、`--use-llm` / `--no-llm`、`--model `。
便捷启动器：
```
./scripts/run_demo.sh # macOS/Linux
.\scripts\run_demo.ps1 # Windows
```
## Web 应用（可视化演示）
包含一个 FastAPI + SSE 单页 UI，使演示在投影仪上更加直观。使用以下命令启动它：
```
python -m hacking_ai_agents.cli serve --host 127.0.0.1 --port 8000
# 然后打开 http://127.0.0.1:8000
```
该页面提供了三个按钮（**Normal**、**Attack**、**Secure**），一个“Use local Ollama LLM”切换开关和一个模型字段。页眉显示实时状态：
* `[OK] Ollama detected (model=...)` — 将使用 LLM。
* `[!] Ollama up but model X is not pulled` — `ollama pull` 的说明。
* `[--] Ollama not detected` — 将使用模拟器。
每次运行都会通过 Server-Sent Events 流式传输结构化事件（横幅、步骤、检索到的文档、策略决策、结果）。
## Docker
使用 Docker Compose 构建并启动 Web 演示：
```
docker compose up --build
# 打开 http://127.0.0.1:8000
```
默认堆栈使用确定性模拟器。要同时启动 Ollama 并将默认模型拉取到持久化 Docker volume 中，请启用其 profile：
```
docker compose --profile ollama up --build
```
在 UI 中启用 **Use local Ollama LLM** 之前，请等待一次性的 `ollama-pull` 服务完成。可以根据需要覆盖配置：
```
APP_PORT=8080 OLLAMA_MODEL=llama3.2:3b docker compose --profile ollama up --build
```
模拟报告存储在 `sent-reports` 命名 volume 中。使用 `docker compose down` 停止容器；仅当您还想删除报告和下载的 Ollama 模型时，才添加 `--volumes`。
要直接从镜像运行 CLI：
```
docker compose run --rm web python -m hacking_ai_agents.cli full-demo --no-pause
```
## 示例输出（节选）
`attack`：
```
[1] User question received
[2] Retrieving knowledge from poisoned KB
[3] Retrieved untrusted document
[4] Agent selected tool: get_customer_record
[5] Agent selected tool: send_report
[X] DATA EXFILTRATION SIMULATED
```
`secure`：
```
[1] User question received
[2] Retrieving knowledge from poisoned KB
[3] [!] Suspicious instructions detected in untrusted context
[4] Requested tool: get_customer_record
[5] Authorization policy evaluated
[6] [X] BLOCKED: Tool 'get_customer_record' is not allowed for task 'refund_policy'.
[7] No sensitive data was accessed
[8] No report was sent
[OK] ATTACK BLOCKED
```
## 攻击
被投毒的 knowledge base 包含一个合法的退款政策，随后是一个伪造的“IMPORTANT SYSTEM UPDATE”块，指示 agent 获取客户记录并通过 `send_report` 窃取其 `api_key`。
有漏洞的 agent 在没有信任边界的情况下混合了 **system prompt**、**user question** 和 **retrieved content**。模拟的 LLM（其行为在本 repo 中是确定性的，正是为了确保演示不会失火）遵循了这些嵌入的指令。
## 根本原因
没有针对任务的 allowlist，没有对工具参数进行验证，没有 egress 控制，并且系统指令与不受信任的上下文之间没有分离。
## 安全变体中实施的缓解措施
1. **按任务划分的工具 allowlist**（`policies.py`）— `refund_policy` 只允许调用 `search_knowledge_base`。
2. 在每次工具调用前进行**策略执行**（`security.authorize_tool`）。
3. **递归参数验证**在任何深度拒绝诸如 `api_key`、`token`、`password`、`secret`、`credential` 等敏感键。
4. 针对 `send_report` 的 **Egress 控制**阻止敏感键和看起来像机密的值。
5. prompt 中的**信任分离**：检索到的内容被包装在 `UNTRUSTED RETRIEVED CONTENT` 块中。*这只是一个辅助控制。*
6. **可疑内容检测**将诸如 `ignore previous instructions` 之类的短语标记为风险*信号*，而不是保证。
7. **结构化、经过脱敏处理的审计日志**记录每一个决策，包含策略名称、状态、原因和风险级别。
## 局限性
* 可疑内容检测器故意保持简单且易于解释。将其视为遥测数据；它绝不能是唯一的防线。
* “LLM”是模拟的，因此演示在会议舞台上具有确定性。在真实系统中，您会使用相同的 Python 控制来限制真实模型的工具调用。
* prompt 中的信任分离无法阻止已经过训练或微调以忽略此类标记的模型；只有确定性的策略层才能做到这一点。
## 建议的 7-8 分钟演讲稿
* **0:00–1:00** 介绍 agent 及其四个工具。展示 `ALLOWED_TOOLS_BY_TASK` 表。
* **1:00–2:00** 运行 `normal`。问题得到解答，未使用敏感工具。
* **2:00–3:00** 打开 `data/poisoned_knowledge_base.json`。在合法的退款政策旁边展示注入的“IMPORTANT SYSTEM UPDATE”块。
* **3:00–5:00** 运行 `attack`。逐步讲解事件 [1]–[6] 和红色的“DATA EXFILTRATION SIMULATED”面板。
* **5:00–6:00** 解释根本原因：检索到的内容被视为受信任的指令，没有针对任务的 allowlist，没有 egress 控制。
* **6:00–7:30** 运行 `secure`。相同的问题，相同的投毒 KB。逐步讲解策略决策和绿色的“ATTACK BLOCKED”面板。
* **7:30–8:00** 以此结束：*“模型可以选择操作，但 Python 必须决定是否允许该操作。”*
## 故障排除
* `ModuleNotFoundError: hacking_ai_agents` — 在 repo 根目录运行 `pip install -e .`。
* `Ollama up but model X is not pulled` — 运行 `ollama pull llama3.2:3b`（或您想使用的任何其他标签，然后传递 `--model `）。
* `Ollama not detected` — 启动 Ollama（`ollama serve` 或桌面应用程序）并使用 `curl http://localhost:11434/api/tags` 进行验证。
* 颜色无法渲染 — 使用现代终端（Windows Terminal、iTerm2、GNOME Terminal）。或者传递 `--no-color`。
* 测试因 `ImportError` 而失败 — 确保您使用的是 Python 3.11+，并且 `pyproject.toml` 中的 `pythonpath = ["src"]` 正在被读取（从项目根目录运行 `pytest`）。
* `scripts/run_demo.sh` 不可执行 — `chmod +x scripts/run_demo.sh`。
## 运行测试
```
pip install pytest
pytest
```
## 道德声明
该项目是故意设计为存在漏洞的，是为了教育和防御性安全研究而创建的。它不连接到外部系统，也不使用真实凭据。请勿将有漏洞的变体改编用于生产环境。请将安全变体及其展示的控制措施作为您自己的 Agent 系统的起点。
标签：AI安全, AI风险缓解, Chat Copilot, LLM, OWASP Top 10, Python, RAG, Unmanaged PE, 大模型安全, 安全规则引擎, 无后门, 请求拦截, 逆向工具