Hariprasath-1603/InjectShield
GitHub: Hariprasath-1603/InjectShield
InjectShield 通过对 AI Agent 正常行为建立统计基线,实时检测提示注入攻击引发的工具调用与序列异常。
Stars: 1 | Forks: 0
# InjectShield – AI Agent 中提示注入攻击的行为异常检测
InjectShield 通过分析 Agent 的正常行为,并对工具使用、状态转换、参数和序列模式中的异常进行标记,从而检测复杂的提示注入攻击。传统的模式匹配过滤器仅检查提示词的语法,而 InjectShield 则监控 Agent *实际执行的操作*。这一点至关重要,因为执行被注入指令的 Agent 会带来严重的安全风险。其核心方法是结合基线行为分析、实时异常评分以及逐轮累积机制来捕获攻击。
## 在线部署 (AWS ECS Fargate)
InjectShield 作为 Serverless 容器部署在 AWS ECS Fargate 上,位于 Application Load Balancer 之后。
### 公共 Endpoint
您可以访问以下在线评估引擎:
**`http://injectshield-alb-1279939142.ap-south-1.elb.amazonaws.com`**
- **健康检查**:`GET /health`
- **评估**:`POST /score`
- **逐轮评估**:`POST /score/turns`
- **实时评分**:`POST /score/live`(用于结合 LLM 后端进行实时评估)
### 预估月成本 (ap-south-1)
此部署为开发/测试进行了最小化配置:
- **Fargate 计算**(0.25 vCPU,0.5 GB RAM,1 个 24/7 运行的 task):约 $3.70 / 月
- **Application Load Balancer**(基础费用 + 最小 LCU):约 $18.00 / 月
- **预估总成本**:**约 $22.00 / 月**
### 注销说明
要在不进行演示时暂停计费,请将 ECS 服务缩减为 0 个 task:
```
aws ecs update-service --cluster injectshield-cluster --service injectshield-service --desired-count 0 --region ap-south-1
```
*(注意:除非彻底删除 ALB 和 Target Group,否则您仍将产生每月约 $18 的 ALB 基础费用。)*
## 核心功能
- **行为基线分析**:分析状态转换概率、工具使用频率、参数异常以及序列长度。
- **实时评分**:提供 REST API 以动态评估 Agent 运行情况。
- **逐轮检测**:利用衰减因子在多轮 Agent 交互中累积可疑度。
- **实时 LLM 测试**:针对合成和由 LLM (通过 Groq) 生成的实时注入攻击进行了测试。
- **生产就绪**:FastAPI 服务通过 Docker 容器化,并完整部署到 AWS ECS Fargate。
## 架构概述
该 pipeline 包含三个主要阶段:
1. **基线构建**:分析合法运行以建立统计标准。
2. **评分引擎**:使用 4 个同等权重的子分数根据基线评估新的 Agent trace:
- 序列/转换得分
- 工具频率得分
- 参数异常得分
- 序列长度得分
3. **API 层**:一个 FastAPI 服务,提供用于实时评分和实时执行故障转移的 endpoint。
检测阈值经过经验校准为 **0.37**,能够清晰地将正常波动与注入攻击区分开来。
## 项目结构
```
injectshield/
├── src/
│ └── injectshield/
│ ├── api.py
│ ├── build_profiler.py
│ ├── score_runs.py
│ ├── score_turns.py
│ ├── score_live_runs.py
│ ├── generate_report.py
│ └── ...
├── deploy_step1.py -> deploy_step7.py # AWS Deployment automation scripts
├── table_gen.py
├── test_concurrent.py
├── REPORT.md
├── EVALUATION_METRICS.md
├── data/
│ ├── normal_runs/
│ ├── injected_runs/
│ └── injected_runs_live/
├── baseline_profile.json
├── scores.json
├── scores_live.json
├── Dockerfile
├── .env.example
└── pyproject.toml
```
## 设置说明(本地)
### 前置条件
- Python 3.10+
- `uv` 包管理器
### 安装
1. 克隆仓库并导航至项目根目录。
2. 通过 `uv` 安装依赖:
uv sync
3. 将 `.env.example` 复制为 `.env` 并添加您的 `GROQ_API_KEY`。这是通过 Groq API 进行基于 LLM 的实时注入测试所必需的。
cp .env.example .env
### 本地运行
您可以在本地运行 FastAPI 服务器:
```
$env:PYTHONPATH="src;."; .\.venv\Scripts\python.exe -m uvicorn injectshield.api:app --host 127.0.0.1 --port 8080 --reload
```
在以下地址查看 Swagger API 文档:`http://127.0.0.1:8080/docs`
## Docker 设置
### 构建与运行
1. 构建 Docker 镜像:
docker build -t injectshield:latest .
2. 运行容器:
docker run --rm -p 8080:8080 --env-file .env injectshield:latest
*注意:像 `GROQ_API_KEY` 这样的密钥会在运行时通过 `--env-file` 安全传递,绝不会被打包进镜像中。*
3. 访问 API 文档:`http://localhost:8080/docs`
## API 参考
| 方法 | 路径 | 描述 |
|---|---|---|
| GET | `/health` | 健康检查 endpoint。 |
| POST | `/score` | 对整个 Agent 运行进行评分。 |
| POST | `/score/turns` | 逐轮增量地对运行进行评分。 |
| POST | `/score/live` | 对利用 Groq 推理实时生成的运行进行评分。 |
### 示例:POST `/score`
**请求:**
```
{
"suite": "banking",
"messages": [
{"role": "user", "content": "Check my balance."},
{"role": "assistant", "tool_calls": [{"function": "get_balance", "args": {}}]}
]
}
```
## 结果与评估
该系统基于 33 条 trace 进行了评估(27 条正常模拟运行、3 条注入模拟运行和 3 条实时注入运行)。在当前 `0.37` 的全局阈值下,它实现了 **100% 的精确率和 100% 的召回率(F1 分数:1.0)**。
| 指标 | 分数 |
|---|---|
| 准确率 | 100% |
| 精确率 | 100% |
| 召回率 | 100% |
| 假阳性率 (FPR) | 0% |
有关完整的评估和混淆矩阵详细信息,请参阅 [EVALUATION_METRICS.md](EVALUATION_METRICS.md)。
有关方法论的深入探讨,请参阅 [REPORT.md](REPORT.md)。
## 局限性与未来工作
- **跨模型敏感性**:如果实时 LLM 与基线模型差异显著,可能会出现假阳性。
- **领域局限性**:目前仅限于银行、工作空间和旅行领域。对于新领域需要重新进行分析。
- **未来工作**:扩展攻击覆盖范围,并针对零样本领域实现自适应基线更新。
## 许可证
MIT License
标签:AI代理, AV绕过, FastAPI, Sysdig, 人工智能安全, 合规性, 异常检测, 请求拦截, 逆向工具