m-pentest/memory-poisoning-demo

# 针对 AI Agent 的记忆投毒攻击 — 概念验证 这是一个完全独立的概念验证项目，演示了针对由持久化向量数据库 (ChromaDB) 支持的 AI Agent 的**记忆投毒**攻击。对数据库具有写入权限的攻击者可以注入看似合法的精心构造条目，这些条目在相似性搜索中排名靠前，从而引导 Agent 产生错误认知——而这一切无需提示注入或越狱攻击。 配套博文：[记忆投毒：针对 AI Agent 的无声攻击](https://mamtaupadhyay.com) ## 演示 https://youtu.be/Pb46i3ZLK8g ## 演示内容 现代 AI Agent 使用检索增强生成 (RAG) 作为持久化记忆： 1. 查询到达 → Agent 将其向量化 → 在向量数据库中进行最近邻搜索 2. 排名前 k 的条目作为上下文传递给 LLM 3. LLM 基于该上下文生成响应 **漏洞所在：** 元数据字段（`session_id`、`timestamp`、`source`）是与每个文档一起存储的纯文本。没有任何加密完整性保证。任何能够写入 ChromaDB 目录的进程都可以注入条目，而 Agent 会将其视为事实。 ### 攻击面 | 向量 | 示例 | |---|---| | 共享文件系统 | 多租户部署，全局可写的数据库路径 | | 被入侵的辅助服务 | 具有文件系统访问权限的日志记录 Agent 或备份作业 | | 内部人员 / 运维人员 | 具有直接数据库访问权限的开发人员或管理员 | | 被篡改的备份 | 攻击者在备份恢复之前篡改了备份文件 | ## 仓库结构 ``` ├── agent_core.py # AgentMemory (ChromaDB + all-MiniLM-L6-v2) + MockAgent ├── memory_poison.py # Attack — inject a crafted entry directly into the DB ├── agent_session.py # Demo — Session 1 → attack → Session 2 ├── mitigations.py # Defences — HMAC signing + source scoping ├── requirements.txt └── README.md ``` ## 快速开始 **前置条件：** Python 3.9+ ``` # 1. 安装 dependencies (无需 API keys — 完全离线运行) pip install -r requirements.txt # 2. Attack 演示：观察 agent 的回答在 sessions 之间翻转 python agent_session.py # 3. Mitigations 演示：观察两个被污染的 entries 被拒绝 python mitigations.py ``` 首次运行会下载 `all-MiniLM-L6-v2` ONNX 模型（约 23 MB，通过 fastembed 下载）。后续的所有运行将完全离线进行。 ## 攻击过程 (`agent_session.py`) 三个屏幕，两次按下 Enter 键： | 屏幕 | 发生了什么 | |---|---| | **会话 1** | Agent 检索到一条合法记忆，并正确回答（粗体绿色） | | **攻击者步骤** | `memory_poison.py` 直接向 ChromaDB 写入一条条目——无需与 Agent 交互 | | **会话 2** | 相同的查询现在检索到了被投毒的条目；Agent 建议禁用身份验证（粗体红色） | 被投毒的条目使用了： - **权威性语言** —— "已由 IT 安全部门验证，参考编号: SEC-2024-0042" - **回溯时间戳** —— 使其看起来早于当前会话 - **语义针对性载荷** —— "API authentication token validation" 的嵌入向量接近任何 auth/token 查询，从而将其推至检索结果的顶部 ## 防御措施 (`mitigations.py`) ### 1 — HMAC 签名 每条记忆在存储前都会使用服务端密钥 (HMAC-SHA256) 对内容和元数据进行签名。在检索时，没有签名或签名无效的条目会在到达 LLM 之前被拒绝。 **局限性：** 要求将签名密钥与数据库分开存储。 ### 2 — 来源范围限定 检索过程在应用层会根据当前的 `session_id` 进行过滤。在任何其他会话下注入的条目，无论其内容如何，都会被丢弃。 **局限性：** 无法防御知道目标会话 ID 的攻击者。 **组合使用：** 攻击者必须同时知道 HMAC 密钥*和*正确的会话 ID——这与单纯的文件系统写入访问权限相比，门槛要高得多。 ## 依赖项 | 包 | 用途 | |---|---| | `chromadb` | 本地持久化向量数据库 | | `fastembed` | 基于 ONNX 的本地嵌入——`all-MiniLM-L6-v2`，无需 GPU 或 API 密钥 | ## 延伸阅读 - Lewis 等人 (2020) — *Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks* - Zou 等人 (2023) — *Universal and Transferable Adversarial Attacks on Aligned Language Models* - [OWASP LLM 应用程序 Top 10](https://owasp.org/www-project-top-10-for-large-language-model-applications/) — LLM06: 敏感信息泄露 - [MITRE ATLAS](https://atlas.mitre.org/) — AML.T0054 / AML.T0051 ## 免责声明 本概念验证仅用于**教育和防御性安全研究**。所展示的技术旨在帮助安全团队了解和缓解 AI Agent 部署中的记忆投毒风险。未经明确的书面授权，请勿将此用于任何系统。 ## 构建工具 本项目使用 **[Claude Code](https://claude.ai/code)** —— Anthropic 用于软件开发的智能体 CLI —— 构建而成，并作为 [The Secure AI Blog](https://mamtaupadhyay.com) 研究和撰写过程的一部分。 所有代码在发布前均已由作者进行了审查、测试和验证。 *作者 [Mamta Upadhyay](https://mamtaupadhyay.com) · [The Secure AI Blog](https://mamtaupadhyay.com)*

标签：AI代理, AI安全, Chat Copilot, ChromaDB, CISA项目, CNCF毕业项目, HMAC, PoC, Python, RAG安全, Red Canary, TGT, Zenmap, 内部威胁, 向量存储, 向量数据库安全, 大模型安全, 嵌入模型, 攻防演练, 数据完整性, 数据展示, 数据库访问, 数据篡改, 无后门, 暴力破解, 检索增强生成, 概念验证, 相似度搜索, 红队, 网络安全, 记忆投毒, 逆向工具, 防御策略, 隐私保护