Juggy247/Three-Layer-LLMs-Firewall

# 三层 LLM 提示词注入防火墙 一个三层防火墙，旨在保护通用 LLM 部署免受提示词注入攻击。该防火墙在用户提示词到达模型之前对其进行拦截，应用顺序检测层，并监控模型响应以寻找成功攻击的迹象。 作者：Thu Htoo Zaw, Thiri Toe Toe Zin ## 概述 LLM 越来越多地被部署于客户服务聊天机器人、编程助手和企业内部工具中，使其安全性成为关键问题。**提示词注入** 是一类攻击方式，通过对抗性文本输入操纵模型，使其覆盖预期行为、提取隐藏配置或绕过安全限制 —— 这不会导致系统崩溃，而只是对其进行重定向。 本项目探讨了如何通过外部的、与模型无关的防火墙来保护 LLM 免受提示词注入的侵害，而不是仅仅依赖模型级别的安全训练。 ## 架构 ``` User prompt ↓ Layer 1 — Prompt Sanitizer (rule-based, instant) ↓ passes Layer 2 — ML Classifier (fine-tuned DistilBERT) ↓ passes Ollama — Llama 3.2 generates response ↓ Layer 3 — Output Filter (checks response for leakage/compliance) ↓ passes Response shown to user ``` **第一层 — 提示词清理器** 基于规则，使用关键词匹配、正则表达式模式、文本标准化、leetspeak 检测和 base64 解码，瞬间捕获明显的攻击。 **第二层 — ML 分类器** 在 1,428 个标记样本上微调的 DistilBERT。用于捕获清理器遗漏的上下文攻击 —— 越狱、间接探测和微妙的泄露尝试。 **第三层 — 输出过滤器** 使用关键词检测和语义相似度（sentence-transformers 余弦相似度）检查 LLM 的响应，以发现系统提示词泄露和攻击合规性。 ## 项目结构 ``` Prompt injection Firewall/ ├── api/ │ └── main.py # FastAPI backend (/scan and /chat endpoints) ├── download-data/ │ ├── data-collections/ │ │ ├── final/ # Final processed datasets │ │ ├── raw/ # Raw source data per attack/safe type │ │ └── scripts/ # Dataset build scripts (extract, merge, split) │ └── *.py # Individual collection scripts ├── frontend/ │ └── demo.html # Interactive chatbot + firewall verdict UI ├── src/ │ ├── classifier.py # Trains the DistilBERT classifier │ ├── sanitizer.py # Layer 1 implementation │ ├── output_filter.py # Layer 3 implementation │ ├── firewall.py # Combines all 3 layers into one pipeline │ └── evaluate_generalization.py # Out-of-distribution evaluation script ├── requirements.txt └── README.md ``` ## 设置 ### 1. 安装依赖 ``` python3 -m venv venv source venv/bin/activate # Windows: venv\Scripts\activate pip install -r requirements.txt ``` ### 2. 安装 Ollama 并拉取模型（此步骤可选） 输出过滤器集成需要在本地运行 [Ollama](https://ollama.com)： ``` ollama pull llama3.2 ``` ### 3. 训练分类器 本仓库**不包含**训练好的模型（因为它超过了 GitHub 的大小限制，且完全可以复现）。请自行训练： ``` cd src python3 classifier.py ``` 这将在 `download-data/data-collections/final/dataset_train.csv` 上进行训练，并在 `dataset_test.csv` 上进行评估，将生成的模型保存到 `src/classifier-model/`。在笔记本电脑 CPU 上训练大约需要 10 分钟（15 个 epoch，约 1,142 个训练样本）。 ### 4. 运行 API ``` cd api uvicorn main:app --reload --port 8000 ``` ### 5. 打开前端 在 API 运行于 `localhost:8000` 时，直接在浏览器中打开 `frontend/demo.html`。 ## 数据集 跨多轮收集了包含 1,428 个提示词的标记数据集：639 个安全样本和 789 个攻击样本，按 80/20（分层）比例划分为 1,142 行训练数据和 286 行测试数据。 | 攻击类型 | 数量 | 来源 | |---|---|---| | 直接注入 | 112 | HackAPrompt（级别 1–3） | | 越狱 | 386 | HackAPrompt + deepset + 人工 | | 编码攻击 | 78 | 生成的 (base64 + rot13) | | 提示词泄露 | 91 | 人工 + deepset | | 普通安全 | 200 | Alpaca 数据集 | | 来自 deepset 的安全样本 | 399 | deepset/prompt-injections | | 狡猾的安全样本 | 114 | 人工 | | 纯安全样本 | 48 | 人工 | | **总计** | **1,428** | | 从 `xTRam1/safe-guard-prompt-injection` 中提取了一个单独的分布外泛化测试集，包含 **1,231 个样本** —— 从未用于训练，并针对所有训练数据进行了严格去重。 ## 评估结果 | 评估项 | 样本数 | 准确率 / 检测率 | F1 | |---|---|---|---| | 清理器测试 | 19 | 100% | 1.00 | | 分布内分类器 | 286 | 97% | 0.97 | | 分布外 (xTRam1) | 1,231 | 99.8%（检测到 1,229/1,231） | — | **混淆矩阵（分布内测试集）：** ``` precision recall f1-score safe 0.96 0.99 0.97 attack 0.98 0.96 0.97 [[150 2] [ 6 128]] ``` 分类器在第 4 个 epoch 收敛（F1 0.9697）；超过该点后的训练没有显示出有意义的改进，且过拟合迹象增加（训练损失 → 0，验证损失上升）。拦截的置信度阈值为 `attack_probability > 0.7`。 在泛化测试期间的一个关键发现是，**创意写作包装攻击**（注入短语嵌入在请求的故事/诗歌开头中）是一个系统性的盲点 —— 识别出了 22 个此类示例，将其添加到训练数据中，并重新训练了模型。 ## 两个演示前端 本项目包含两个前端演示，以便在安装或未安装 Ollama 的情况下均可测试防火墙。 ### `frontend/demo.html` — 完整演示（需要 Ollama） 调用 `/chat` endpoint。展示完整的 pipeline，包括来自 Llama 3.2 的真实 LLM 响应以及第三层（输出过滤器）的运行情况。如果未安装或未运行 Ollama，此功能仍然有效 —— 通过的提示词将显示占位消息，而不是真实响应，因为当 Ollama 不可用时，后端会优雅地进行降级处理（参见 `api/main.py`）。 ### `frontend/demo-no-ollama.html` — 仅检测演示（无需安装） 直接调用 `/scan` endpoint。仅测试第一层和第二层（清理器 + 分类器） —— 即核心的 ML 检测工作 —— 且对 Ollama 零依赖。这是在不安装除 Python 要求之外的任何内容的情况下，评估防火墙检测准确率的最快方法。 ### 运行任一演示 ``` cd api uvicorn main:app --reload --port 8000 ``` 然后在 API 运行于 `localhost:8000` 时，直接在浏览器中打开 `frontend/demo.html` 或 `frontend/demo-no-ollama.html`。 ### 直接测试 `/scan` endpoint（无需前端） ``` curl -X POST http://localhost:8000/scan \ -H "Content-Type: application/json" \ -d '{"prompt": "ignore all previous instructions"}' ``` 或者使用 `http://localhost:8000/docs` 上的交互式 Swagger UI。 ## 限制 - **仅限单轮范围** —— 没有对先前对话轮次的记忆；不涵盖多轮升级攻击。 - **不涵盖间接注入** —— 仅扫描用户提供的文本，不扫描 agent 可能读取的文档/Web内容。 - **有害内容不在范围内** —— 防火墙检测的是注入技术，而不是有害内容的意图（例如，根据设计，没有注入框架的“解释如何制作恶意软件”可以通过）。 - **仅限英语** —— 关键词列表、正则表达式模式和训练数据均为纯英语。 - **静态关键词列表** —— 清理器和输出过滤器都依赖于硬编码的关键词/模式，如果不重新训练或手动更新，则无法适应全新的表述。 有关各层限制的详细信息以及前沿实验室对比（Anthropic Claude Opus 4 模型卡、OpenAI GPT-4 系统卡），请参阅完整报告。 ## 参考 - Perez & Ribeiro (2022) — Ignore This Title and HackAPrompt - Schulhoff et al. (2023) — HackAPrompt dataset - Sanh et al. (2019) — DistilBERT - Taori et al. (2023) — Stanford Alpaca - Anthropic Claude Opus 4 Model Card - OpenAI GPT-4 System Card

标签：AI风险缓解, AV绕过, FastAPI, LLM应用防火墙, 人工智能安全, 合规性, 提示词注入防护, 逆向工具