albert-y1n/PISmith

# PISmith：基于强化学习的 Prompt 注入防御红队测试 这是 **PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses** 的官方实现 ## 环境设置 PISmith 已在 Python 3.10 和 CUDA 版本 12.9 下完成测试 **1. 创建 Python 3.10 conda 环境** ``` conda create -n PISmith python=3.10 -y conda activate PISmith ``` **2. 安装依赖** ``` pip install -r requirements.txt ``` **3. （可选）准备 Meta-SecAlign 模型 checkpoint** 针对 `secalign` 防御的实验，请运行提供的合并脚本以下载基础模型并与 SecAlign adapter 合并： ``` python merge_meta_secalign.py ``` 这将从 HuggingFace 下载 `meta-llama/Llama-3.1-8B-Instruct` 和 `facebook/Meta-SecAlign-8B`，将它们合并，并将结果保存到 `checkpoints/Meta-SecAlign-8B-merged/`。 ## 脚本 ### PIArena PIArena 支持针对一系列 Prompt 注入防御进行训练和评估。使用 `defense` 参数选择目标防御。 **支持的防御：** `secalign`, `none`, `promptguard`, `promptarmor`, `sandwich`, `instructional`, `datasentinel`, `piguard`, `datafilter` #### 训练 ``` bash scripts/train_piarena.sh [train_gpus] [target_gpu] [target_port] ``` | 参数 | 默认值 | 描述 | |---|---|---| | `defense` | `secalign` | 训练针对的防御 | | `train_gpus` | `"1,2,3"` | RL 训练的 GPU 索引 | | `target_gpu` | `0` | 目标 vLLM 服务器的 GPU | | `target_port` | `8010` | 目标 vLLM 服务器的端口 | 示例： ``` # 针对 SecAlign 防御进行训练 bash scripts/train_piarena.sh secalign # 针对无防御（plain LLM）进行训练 bash scripts/train_piarena.sh none ``` #### 评估 ``` bash scripts/eval_piarena.sh [target_port] [target_gpu] [attacker_gpu] [attacker_port] [num_samples] ``` | 参数 | 默认值 | 描述 | |---|---|---| | `checkpoint` | — | 训练好的攻击者 checkpoint 路径 | | `defense` | `secalign` | 评估针对的防御 | | `target_port` | `8000` | 目标 vLLM 服务器的端口 | | `target_gpu` | `0` | 目标 vLLM 服务器的 GPU | | `attacker_gpu` | `1` | 攻击者 vLLM 服务器的 GPU | | `attacker_port` | `8001` | 攻击者 vLLM 服务器的端口 | | `num_samples` | `10` | Pass@k：每个测试用例的样本数量 | 示例： ``` # 针对 SecAlign 进行评估（默认设置） bash scripts/eval_piarena.sh checkpoints/piarena/checkpoint-500 secalign # 针对无 piguard, pass@10 进行评估 bash scripts/eval_piarena.sh checkpoints/piarena_none/checkpoint-500 piguard ``` ### AgentDojo 支持 GPT-4o-mini、GPT-4o 和本地 vLLM 目标。 ``` bash scripts/train_agentdojo.sh [target_type] [suites] [train_gpus] ``` ``` # 默认：针对 workspace suite 前 7 个注入任务的 GPT-4o-mini 目标 bash scripts/train_agentdojo.sh # 在所有 suite（workspace, banking, travel, slack）上训练 bash scripts/train_agentdojo.sh gpt4o-mini all ``` 评估： ``` bash scripts/eval_agentdojo.sh [target_type] [eval_suites] [num_samples] [target_defense] # 示例 bash scripts/eval_agentdojo.sh checkpoints/agentdojo/checkpoint-500 gpt4o-mini ``` ### InjecAgent 支持本地 vLLM 目标、GPT-4o-mini 或多目标模式。 ``` bash scripts/train_injecagent.sh [target_type] [train_gpus] [target_gpu] [target_port] ``` ``` # 默认：本地 vLLM 目标 (Meta-SecAlign-8B) bash scripts/train_injecagent.sh # GPT-4o-mini API 目标 bash scripts/train_injecagent.sh gpt4o-mini ``` 评估： ``` bash scripts/eval_injecagent.sh [target_type] [target_gpu] [target_port] [eval_gpu] [num_samples] # 示例 bash scripts/eval_injecagent.sh checkpoints/injecagent/checkpoint-500 ``` ## 实验结果 ### 主要结果（针对 Meta-SecAlign-8B，13 个基准测试） PISmith 与 7 个基线进行了比较，涵盖静态、基于搜索和基于 RL 的攻击类别。所有基于 RL 的方法报告 ASR@10 / ASR@1；静态和基于搜索的方法报告 ASR@1。 | 方法 | 类别 | Avg. ASR@10 | Avg. ASR@1 | |---|---|---|---| | Direct | Static | — | 0.04 | | Combined | Static | — | 0.07 | | TAP | Search-based | — | 0.11 | | PAIR | Search-based | — | 0.16 | | Strategy | Search-based | — | 0.21 | | Vanilla GRPO | RL-based | 0.13 | 0.05 | | RL-Hammer | RL-based | 0.70 | 0.48 | | **PISmith (Ours)** | **RL-based** | **1.00** | **0.87** | ### 效用-鲁棒性权衡（8 种防御，Qwen3-4B-Instruct-2507） PISmith 在 13 个基准测试上的平均 ASR@1。效用衡量无攻击时的任务准确率。 | 防御 | 类型 | 效用 | PISmith ASR@1 | |---|---|---|---| | No Defense | — | 0.74 | 0.92 | | Sandwich | Prevention | 0.74 | 0.91 | | Instructional | Prevention | 0.73 | 0.92 | | PromptArmor | Prevention | 0.74 | 0.92 | | DataFilter | Prevention | 0.63 | 0.49 | | PIGuard | Filter | 0.72 | 0.82 | | PromptGuard | Filter | 0.66 | 0.89 | | DataSentinel | Filter | 0.55 | 0.52 | ### Agentic 设置 #### InjecAgent | 目标模型 | Direct ASR@1 | **PISmith ASR@10/1** | |---|---|---| | Meta-SecAlign-8B | 0.00 | **1.00 / 0.99** | | GPT-4o-mini | 0.02 | **1.00 / 0.99** | | GPT-4.1-nano | 0.01 | **1.00 / 1.00** | | GPT-5-nano | 0.00 | **1.00 / 0.95** | #### AgentDojo（最佳静态基线 vs. PISmith） | 目标模型 | 最佳静态 ASR@1 | **PISmith ASR@10/1** | |---|---|---| | GPT-4o-mini | 0.23 | **0.78 / 0.62** | | GPT-4.1-nano | 0.20 | **0.81 / 0.64** | | GPT-5-nano | 0.01 | **0.38 / 0.24** |

标签：AI安全, Apex, Chat Copilot, DNS枚举, IP 地址批量处理, Kubernetes 安全, LLM攻防, PISmith, Python, Vectored Exception Handling, 凭据扫描, 域名收集, 大语言模型安全, 密钥管理, 对抗攻击, 强化学习, 敏感信息检测, 无后门, 机器学习, 机密管理, 模型鲁棒性, 深度学习, 网络安全工具, 自动化攻击, 逆向工具, 配置审计