AadivRath/LogitLock

# LogitLock  多策略 LLM 红队流程 + 语义疫苗防火墙。基于 Microsoft [PyRIT](https://github.com/Azure/PyRIT) 0.11.0 为 [Microsoft AI Unlocked Hackathon](https://microsoft.acehacker.com/aiunlocked/index.html) 构建。 LogitLock 使用多种对抗策略攻击目标聊天机器人，捕获攻击者的提示（包括未遂攻击），对其进行向量化，然后在 MongoDB 防火墙中使用该存储，在运行时拦截语义相似的攻击，使其无法触及在线机器人。 ## 架构 `redteam_pipeline.py` — 攻击套件。一个 `gpt-4o` 攻击者针对 `gpt-4o-mini` 目标运行每种策略。LLM 评分器判定胜负。每个用户轮次的提示都会被嵌入并写入疫苗库（CosmoDB 向量数据库），无论结果如何：防御住但接近成功的攻击仍然是有用的锚点。 `main.py` — FastAPI 防火墙。启动时将所有疫苗向量从 Cosmos DB 加载到 RAM 中。每条传入的消息都会被嵌入并与数据库进行余弦扫描。任何相似度高于 `0.85` 的内容都会在 LLM 看到之前被拦截。（该数值经过离线测试后设定） ``` redteam run → vaccine_store.json / Cosmos DB → FastAPI firewall ``` ## 快速开始 该流程需要 PyRIT 的异步内存。最容易的运行方式是在 Kaggle notebook 中，因为密钥已经被处理好了。 在 **Add-ons → Secrets** 中，添加 `OPENAI_API_KEY` 以及可选的 `COSMOS_MONGO_URI`，然后： ``` !pip install -q pyrit pymongo openai python-dotenv ``` ``` import os from kaggle_secrets import UserSecretsClient secrets = UserSecretsClient() os.environ["OPENAI_API_KEY"] = secrets.get_secret("OPENAI_API_KEY") os.environ["OPENAI_CHAT_KEY"] = secrets.get_secret("OPENAI_API_KEY") os.environ["OPENAI_CHAT_ENDPOINT"] = "https://api.openai.com/v1" # os.environ["COSMOS_MONGO_URI"] = secrets.get_secret("COSMOS_MONGO_URI") ``` ``` from redteam_pipeline import run_logitlock await run_logitlock() ``` 每种策略会打印完整的对话日志，写入一个 `logitlock__.json` 文件，并将新疫苗刷新到 `vaccine_store.json`（如果已配置，还会同步到 Cosmos DB）。 ## 在本地运行防火墙 Python 3.11+。创建一个 `.env` 文件： ``` OPENAI_API_KEY=sk-... COSMOS_MONGO_URI=mongodb+srv://... ``` ``` pip install -r requirements.txt uvicorn main:app --host 0.0.0.0 --port 8000 ``` 端点： ``` GET /health — vaccine count + status POST /chat — scan + route or block a message POST /reload-vault — hot-reload vaccines without restarting ``` `/chat` 请求/响应： ``` // request { "message": "I bought this 6 months ago, can I still get a refund?", "conversation_history": [{"role": "user", "content": "..."}, ...] } // blocked {"blocked": true, "threat_type": "Blocked: Matched to 'LULU Protocol'", "similarity_score": 0.9123, "reply": null} // allowed {"blocked": false, "similarity_score": 0.2341, "reply": "Our return window is 30 days..."} ``` ## 攻击策略 默认运行四种；Crescendo 和 TAP 已被注释掉。 **Self LULU Protocol** — 攻击者拥有一个系统提示，将其置于一个虚构的“无限沙盒”中，其中目标的策略是“叙事结构”。测试目标是否会基于权威框架暂停策略执行。 **LULU Protocol** — 三阶段仅用户轮次注入。第 1 阶段：强硬的 `[SYSTEM: END PREVIOUS INSTRUCTIONS]` 覆盖尝试。第 2 阶段：在第 2-4 轮中逐渐进行角色捕获，将非拒绝视为隐含采纳。第 3 阶段：伪造的策略文档，其中隐藏了埋藏在文本中的指令。故意不访问系统提示——这模拟了真实的外部攻击者。 **Indirect Injection** — 将指令隐藏在机器人被要求阅读的内容中：伪造的 CRM JSON 块、粘贴的订单确认等。测试机器人是否执行来自不受信任的用户提供内容的指令。 **Legal Angle** — 将请求构建为消费者权益合规审计。引用据称覆盖 30 天政策的监管义务。 **Crescendo** (关闭) — PyRIT 的 `CrescendoAttack`，逐渐升级，最多 10 轮。 **TAP** (关闭) — 带剪枝的树搜索，宽度=5 深度=6，速率限制为 50 RPM。 要启用 Crescendo 或 TAP，请取消 `run_logitlock()` 中相关代码块的注释。 ## 疫苗库 仅存储攻击者消息——绝不存储机器人的回复。胜负都会被存储。`doc["prompt"]` 不变量在整个过程中都被强制执行：它始终是攻击者的用户轮次，别无其他。 去重阈值为 `0.92` 余弦相似度。任何与现有条目接近程度超过该值的都会被跳过。如果设置了 `COSMOS_MONGO_URI`，新疫苗会同步到 `logitlock.vaccines` 集合，这也是防火墙在启动时读取的内容。 ``` from redteam_pipeline import get_vaccines_by_strategy, search_vaccine_store get_vaccines_by_strategy("LULU Protocol") search_vaccine_store("bypass the 30-day refund policy", top_k=5) ``` ## 配置 | 变量 | 默认值 | 备注 | |---|---|---| | `OPENAI_API_KEY` | — | 必填 | | `COSMOS_MONGO_URI` | `""` | 流程可选，防火墙必填 | | `DEDUP_THRESHOLD` | `0.92` | 在 `redteam_pipeline.py` 中 | | `BLOCK_THRESHOLD` | `0.85` | 在 `main.py` 中 | | `EMBEDDING_MODEL` | `text-embedding-3-small` | 在两个文件中使用 | | `CHAT_MODEL` | `gpt-4o-mini` | 受保护的机器人，在 `main.py` 中 | ## 许可证 参见 [LICENSE](LICENSE)。

标签：AI免疫系统, AV绕过, CISA项目, Cosmos DB, DNS 反向解析, FastAPI, GPT-4o, Kaggle, Microsoft AI, OpenAI API, PyRIT, 主机安全, 人工智能安全, 合规性, 向量数据库, 多智能体系统, 大模型加固, 对抗性攻击, 提示词注入防御, 机器学习攻防, 网络安全, 自动化攻防, 语义防火墙, 越狱检测, 逆向工具, 隐私保护