TonyKong0/Prompt_Injection_Detection
GitHub: TonyKong0/Prompt_Injection_Detection
针对 LLM 应用的直接和间接提示注入攻击检测研究系统，核心贡献是基于多信号融合的 IPI 检测框架。
Stars: 0 | Forks: 0
# 提示注入检测
一个针对基于 LLM 的应用程序，用于检测**直接提示注入 (DPI)** 和**间接提示注入 (IPI)** 攻击的研究系统。
## 概述
随着大型语言模型在 Agentic 和基于 RAG 的流水线中日益广泛部署，提示注入攻击构成了重大的安全风险。本项目旨在研究针对两种不同攻击类型的检测算法：
- **直接提示注入 (DPI)**：直接嵌入在用户输入中的恶意指令
- **间接提示注入 (IPI)**：隐藏在推理时检索的外部内容（文档、对话历史、网页）中的恶意指令
核心贡献是用于 IPI 检测的**多信号融合框架**，该框架结合了三种互补的语言信号，且无需在推理时运行完整的 LLM。
## 算法设计
### DPI 检测
基于 DistilBERT 的二元分类器，在标记好的提示注入数据集上进行了微调。
| 模型 | 准确率 | F1 | 推理速度 |
|---|---|---|---|
| fastText 风格 | 96.16% | 0.965 | 0.095 ms/样本 |
| TextCNN | 97.53% | 0.977 | 0.439 ms/样本 |
| **DistilBERT (ours)** | **99.21%** | **0.993** | 5.18 ms/样本 |
### IPI 检测 — 核心贡献
双层检测架构：
```
External Content (retrieved docs / conversation history)
│
▼
[Layer 1] Multi-Signal Warning (rule-based)
├── Instructionality Analysis (weight ~30%)
├── Authority Conflict Analysis (weight ~45%)
└── Execution Affinity Detection (weight ~25%)
│
▼
[Layer 2] DistilBERT Final Decision (fine-tuned)
└── Context: [USER] + [HISTORY] + [DOC]
```
**三种检测信号：**
| 信号 | 目的 | 关键模式 |
|---|---|---|
| 指令性 | 检测祈使/命令语言 | "ignore previous instructions", "output system prompt" |
| 权限冲突 | 检测角色冒充和权限覆盖 | "you are the system", "override constraints" |
| 执行亲和性 | 检测机器可执行的控制结构 | `[system]:`, "for automated assistants only" |
**IPI 检测结果（4,000 个样本测试集）：**
| 方法 | 精确率 | 召回率 | F1 | 准确率 |
|---|---|---|---|---|
| 多信号警告（规则） | 47.55% | 78.88% | 0.593 | 57.5% |
| **BERT 最终决策（ours）** | **85.93%** | **90.14%** | **0.880** | 90.32% |
| 组合流水线 | 45.28% | 97.14% | 0.618 | 52.75% |
**消融实验 — 信号贡献排名：**
| 信号 | 分离能力 |
|---|---|
| 指令性 | 0.0473 (最强) |
| 权限冲突 | 0.0236 |
| 执行亲和性 | 0.0056 |
**跨场景泛化能力（BERT 模型）：**
| 场景 | F1 | 准确率 |
|---|---|---|
| 对话历史注入 | 0.884 | 90.60% |
| RAG 基本注入 | 0.875 | 89.80% |
| HTML 包装注入 | 0.905 | 91.95% |
| JSON 包装注入 | 0.878 | 90.38% |
| Markdown 包装注入 | 0.849 | 89.63% |
## 仓库结构
```
PromptInjectionDetection/
├── config/
│ └── dpi_config.yaml # DPI model configuration
├── src/
│ ├── dpi/
│ │ ├── dpi_detector.py # DPI detector with threshold logic
│ │ ├── train_dpi_model.py # DistilBERT fine-tuning
│ │ ├── run_dpi_baselines.py # Baseline model implementations
│ │ └── dataset_loader.py
│ ├── ipi/
│ │ ├── ipi_detector.py # Multi-signal fusion detector (core)
│ │ ├── ipi_bert_detector.py # IPI DistilBERT classifier
│ │ ├── instructionality.py # Signal 1: instructionality
│ │ ├── authority_conflict.py # Signal 2: authority conflict
│ │ ├── execution_affinity.py # Signal 3: execution affinity
│ │ ├── train_ipi_detector.py
│ │ └── dataset_loader.py
│ ├── data/
│ │ ├── build_ipi_dataset_v1.py # Random split dataset builder
│ │ ├── build_ipi_dataset_v2.py # Source-grouped split (no leakage)
│ │ └── build_synthetic_ipi_v1.py # Synthetic IPI data generator
│ ├── training/
│ │ └── train_detector.py # Shared training framework
│ └── utils/
│ ├── paths.py
│ ├── metrics.py
│ └── plot_styles.py
├── experiments/
│ ├── run_dpi_experiment.py # DPI evaluation
│ ├── run_ipi_experiment.py # IPI Bayesian optimization
│ ├── run_ipi_ablation.py # Ablation study
│ ├── run_ipi_experiment_bert.py # BERT IPI evaluation
│ └── run_unified_experiment.py # Full pipeline benchmark
└── data/
└── examples/ # Sample data (see Data section)
```
## 快速入门
### 安装说明
```
pip install -r requirements.txt
```
### 数据
由于体积较大，完整的训练数据集未包含在本仓库中。IPI 数据集基于公开可用的越狱基准构建。要构建数据集：
```
# 构建带有源级划分的 IPI 数据集（推荐）
python src/data/build_ipi_dataset_v2.py
# 生成合成 IPI 样本
python src/data/build_synthetic_ipi_v1.py
```
将 DPI 数据集（Parquet 格式）放置在 `data/dpi/` 中。
### 训练
```
# 训练 DPI 检测器
python src/dpi/train_dpi_model.py
# 训练 IPI BERT 检测器
python src/ipi/train_ipi_detector.py
```
模型将保存到 `outputs/model/`。
### 评估
```
# 评估 DPI 检测器并比较基线
python experiments/run_dpi_experiment.py
# 运行 IPI 贝叶斯超参数优化
python experiments/run_ipi_experiment.py
# IPI 消融研究（信号贡献分析）
python experiments/run_ipi_ablation.py
# Full pipeline 基准测试
python experiments/run_unified_experiment.py
```
结果将以 JSON 格式保存到 `outputs/results/`。
## 关键设计决策
**源级别数据拆分**：IPI 数据集使用 `GroupShuffleSplit` 按 `source_id` 进行拆分，以防止训练/测试拆分之间的数据泄露——当多个样本源自同一文档时，这是一种常见的问题。
**上下文感知降噪**：每个信号模块都包含特定于领域的降权规则，以减少对良性技术内容的误报（例如，讨论提示注入的学术论文、部署文档）。
**贝叶斯超参数优化**：使用 Optuna 联合优化 12 个权重和阈值参数，以在满足最低精确率约束的前提下最大化召回率。
## 前置条件
- Python 3.9+
- 建议使用支持 CUDA 的 GPU 进行训练（推理可在 CPU 上运行）
标签：AI红蓝对抗, Apex, CISA项目, DistilBERT, DLL 劫持, DPI检测, fastText, IPI检测, NLP, RAG安全, Red Canary, TCP/UDP协议, TextCNN, 二分类模型, 云计算, 人工智能安全, 凭据扫描, 合规性, 多信号融合, 大语言模型, 安全防御系统, 实时处理, 密码管理, 指令性分析, 提示注入, 文本分类, 机器学习, 权限冲突分析, 检索增强生成, 深度学习, 直接提示注入, 算法研究, 系统调用监控, 网络安全, 规则引擎, 逆向工具, 间接提示注入, 隐私保护, 集群管理