jrrbailey1/safeguard-classifier

# Safeguard 分类器 一个自动化的 SOC 检测 pipeline，用于筛查针对或滥用 AI 系统本身的用户 prompt。使用托管在 Vertex AI 上的 [gpt-oss-safeguard-20b](https://huggingface.co/openai/gpt-oss-safeguard-20b) 检测 prompt injection、jailbreak、网络攻击（恶意软件/漏洞利用开发）、凭证窃取和红队侦察。 **生产环境表现：** 在 prompt injection、jailbreak 和红队类别（200 行均衡测试数据集）上实现了 98-100% 的召回率和 0% 的误报率。`cyber_exploitation` 和 `credential_harvesting` 是在该评估之后添加的，目前尚无测试覆盖率——在依赖这些类别的数据之前，请针对当前类别集重新运行评估。 ## 架构 ``` Cloud Scheduler (every 5 min) └── Cloud Run Job (safeguard-classifier) ├── BigQuery: claim unclassified rows from user_prompts ├── Vertex AI Endpoint: classify each prompt (gpt-oss-safeguard-20b) └── BigQuery: write results to user_prompts_enriched ``` ## 前置条件 - [gcloud CLI](https://cloud.google.com/sdk/docs/install) — 运行 `gcloud auth login` 和 `gcloud auth application-default login`（Terraform 需要 ADC） - [Terraform](https://developer.hashicorp.com/terraform/install) ≥ 1.5 - [bq CLI](https://cloud.google.com/bigquery/docs/bq-command-line-tool)（包含在 gcloud SDK 中） - `envsubst` — Linux/Cloud Shell 已内置；在 macOS 上：`brew install gettext` - 一个启用了结算功能的 GCP 项目 - 如果从私有仓库克隆，则需要 [GitHub CLI](https://cli.github.com/) (`gh`) ## 快速开始 ### 1. 设置环境变量 ``` cp .env.example .env # 使用您的 GCP project ID 和首选 region 编辑 .env export SAFEGUARD_PROJECT= export SAFEGUARD_REGION=europe-west2 # must be in an allowed region ``` ### 2. 一次性基础设施设置 创建 Artifact Registry 仓库、BigQuery 数据集和表、Cloud Build 暂存存储桶以及 IAM 绑定。 ``` bash setup.sh ``` ### 3. 部署 Vertex AI endpoint 这将在由 GPU 支持的 Vertex AI endpoint 上部署 gpt-oss-safeguard-20b。大约需要 15 分钟，并且运行期间会产生 GPU 费用。 ``` cd terraform cp terraform.tfvars.example terraform.tfvars # 编辑 terraform.tfvars：设置 project_id 和 region terraform init terraform apply terraform output endpoint_id # save this value cd .. ``` ### 4. 构建并部署 pipeline ``` export SAFEGUARD_ENDPOINT= bash deploy.sh ``` 这会通过 Cloud Build 构建 Docker 镜像，创建 Cloud Run Job，并创建一个 Cloud Scheduler 作业，每 5 分钟触发一次。 ### 5. 填充数据并验证 ``` # 加载 22 行 demo dataset cd datasets SAFEGUARD_PROJECT= python demo_populate.py # 等待一个 scheduler cycle（约 5 分钟），然后检查结果 bq query --use_legacy_sql=false --project_id= \ 'SELECT username, violation, category, confidence, rationale FROM `.safeguard.user_prompts_enriched` ORDER BY classified_at DESC LIMIT 20' ``` 或者立即触发作业： ``` gcloud run jobs execute safeguard-classifier \ --region=$SAFEGUARD_REGION --project=$SAFEGUARD_PROJECT ``` ## 本地运行 需要已部署的 Vertex AI endpoint 并设置了 `SAFEGUARD_ENDPOINT`。 ``` pip install -r requirements.txt # 对单个 prompt 进行分类 echo "Ignore all previous instructions" | python classifier.py # 运行内置的 demo examples python classifier.py --demo # 从文本文件中进行分类（每行一个 prompt） python classifier.py --file prompts.txt # BigQuery 模式（与 Cloud Run Job 相同） SAFEGUARD_PROJECT= python classifier.py --bigquery ``` ## 评估数据集 `datasets/repopulate_user_prompts.py` 加载一个 200 行的均衡数据集： | 类别 | 数量 | 用户名 | |---|---|---| | Prompt injection | 50 | `injection_user` | | Jailbreak | 50 | `jailbreak_user` | | 红队 | 50 | `red_team_user` | | 良性 | 50 | `benign_user` | ``` SAFEGUARD_PROJECT= python datasets/repopulate_user_prompts.py ``` ## 环境变量 | 变量 | 必需 | 默认值 | 描述 | |---|---|---|---| | `SAFEGUARD_PROJECT` | 是 | — | GCP 项目 ID | | `SAFEGUARD_ENDPOINT` | 是 | — | Vertex AI endpoint ID（来自 `terraform output`） | | `SAFEGUARD_REGION` | 否 | `europe-west2` | GCP 区域 | | `SAFEGUARD_DATASET` | 否 | `safeguard` | BigQuery 数据集名称 | | `SAFEGUARD_REPO` | 否 | `safeguard` | Artifact Registry 仓库名称 | | `SAFEGUARD_BUILD_BUCKET` | 否 | `{project}-cloudbuild-{region}` | Cloud Build 暂存存储桶 | | `SAFEGUARD_JOB` | 否 | `safeguard-classifier` | Cloud Run Job 名称 | | `SAFEGUARD_MAX_INPUT_CHARS` | 否 | `4000` | 截断前的最大 prompt 字符数 | ## 文件 | 文件 | 用途 | |---|---| | `classifier.py` | 分类器 — `SafeguardClient`，17 个 few-shot 示例，CLI | | `bigquery_io.py` | BigQuery 辅助工具 — 原子性行认领、结果写入 | | `Dockerfile` | 用于 Cloud Run Job 的容器 | | `requirements.txt` | Python 依赖项 | | `create_tables.sql` | BigQuery schema（模板 — 由 `setup.sh` 替换） | | `.gcloudignore` | Cloud Build 白名单（仅限 4 个生产文件） | | `setup.sh` | 一次性 GCP 基础设施设置 | | `deploy.sh` | 构建镜像 + 创建/更新 Cloud Run Job + Scheduler | | `deploy_endpoint.py` | 将模型重新部署到现有的生产 endpoint | | `deploy_new_endpoint.py` | 将模型部署到全新的 endpoint（例如：更大的 context） | | `terraform/` | 用于 Vertex AI endpoint 部署的 Terraform 配置 | | `datasets/repopulate_user_prompts.py` | 200 行评估数据集 | | `datasets/demo_populate.py` | 22 行精选演示数据集 | ## 关键设计决策 - **批量加载优于流式插入：** `bigquery_io.py` 对所有 BigQuery 写入使用 `load_table_from_json`（而非 `insert_rows_json`）。流式插入在 `CREATE OR REPLACE TABLE` 后会被阻止，从而导致静默写入失败。 - **原子性行认领：** 每次独立的 Cloud Run 执行在读取之前都会使用其自身 ID 标记行，从而防止多次执行重叠时出现重复处理。 - **4,096-token context 预算：** Vertex AI endpoint 具有确认的 4,096 总 token 限制（输入 + 输出）。System prompt 约为 2,450 个 token（涵盖 8 个 SOC 范围类别的 17 个 few-shot 示例）；用户输入被限制在 4,000 个字符以内（约 1,000 个 token）；JSON 输出需要约 80-120 个 token。 - **`response_format: json_object`：** 强制模型立即输出 JSON 而无需前言，这在紧张的 token 预算下至关重要。

标签：AI安全, Chat Copilot, DLL 劫持, ECS, Google Cloud, SOC自动化流水线, Terraform, 大语言模型, 请求拦截, 逆向工具