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, 人工智能安全, 合规性, 异常检测, 请求拦截, 逆向工具