bch1212/queryshield
GitHub: bch1212/queryshield
QueryShield 是一个位于 AI agent 与企业数据库之间的安全 SQL 代理，通过 AST 校验、行级安全和审计日志保障数据访问的合规与安全。
Stars: 1 | Forks: 0
# QueryShield
[](https://github.com/bch1212/queryshield/actions/workflows/test.yml)
AI agent 与企业数据库之间的安全 SQL 代理。
Agent 以纯英文（或结构化 SQL）调用单个 endpoint。QueryShield 会：
1. 通过 Claude 结合 prompt 缓存，将自然语言翻译为 → SQL。
2. 在 AST 级别验证每一条查询 —— 仅允许 `SELECT`，不允许
堆叠语句、禁用函数，且必须包含 LIMIT。
3. 应用针对每个 agent 的行级安全控制：schema/table 白名单以及
`WHERE` 子句注入。
4. 针对客户 DB 执行查询并返回数据行。
5. 将每次尝试记录到只允许追加的审计表中 —— 仅记录元数据，
绝不记录行内容。
Agent 永远不会看到连接字符串。
## 快速开始
```
pip install -r requirements.txt
cp .env.example .env
# 设置 ANTHROPIC_API_KEY、DATABASE_URL、VAULT_KEY（见下文）
python -m queryshield.start
```
为 `VAULT_KEY` 生成一次 Fernet 密钥，并且切勿丢失：
```
python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"
```
## 端到端流程 (curl)
```
# 1) 启动一个 tenant。返回 admin API key — 请将其复制。
curl -X POST localhost:8000/v1/tenants?name=Acme
# 2) 注册客户 DB。连接字符串在静止状态下是加密的。
curl -X POST localhost:8000/v1/databases \
-H 'X-Admin-Key: qs_...' \
-H 'Content-Type: application/json' \
-d '{
"alias": "prod",
"db_type": "postgresql",
"connection_string": "postgresql://reader:secret@db.acme.internal:5432/app",
"allowed_tables": ["users", "orders"]
}'
# 3) 为你的 AI app 提供一个 scoped agent（不同于 admin）。
curl -X POST localhost:8000/v1/agents \
-H 'X-Admin-Key: qs_...' \
-H 'Content-Type: application/json' \
-d '{ "name": "reporting", "tenant_id": "" }'
# 4) 设置该 agent 的 RLS 策略。
curl -X POST localhost:8000/v1/policies \
-H 'X-Admin-Key: qs_...' \
-H 'Content-Type: application/json' \
-d '{
"agent_id": "",
"database_alias": "prod",
"allowed_tables": ["users", "orders"],
"row_filters": { "users": "tenant_id = 42" }
}'
# 5) 该 agent 进行查询。
curl -X POST localhost:8000/v1/query \
-H 'X-API-Key: qs_...' \
-H 'Content-Type: application/json' \
-d '{
"database_alias": "prod",
"query": "how many active users do we have?",
"mode": "nl",
"max_rows": 10
}'
```
## MCP 集成
**已在[官方 MCP Registry](https://registry.modelcontextprotocol.io/v0/servers?search=queryshield) 中收录，名称为 `io.github.bch1212/queryshield`。**
安装客户端：
```
pip install queryshield-mcp
```
然后将以下内容放入你的 Claude Desktop / Cursor / agent 配置中：
```
{
"queryshield": {
"command": "queryshield-mcp",
"env": { "QUERYSHIELD_API_KEY": "qs_..." }
}
}
```
独立的 PyPI package 源码位于 `packages/queryshield-mcp/`。
对于 Glama 等 MCP 目录评估器，仓库根目录还包含一个轻量级的 `Dockerfile`，用于启动已发布的 `queryshield-mcp` stdio server 以进行工具内省。该容器在 MCP 初始化/工具发现阶段不需要 `QUERYSHIELD_API_KEY`；仅当实际调用发现的工具以访问 QueryShield API 租户时，才需要此密钥。
## MCP 集成（旧版）
将以下内容放入任何支持 MCP 的客户端中（Claude Desktop、Cursor、自定义 agent）：
```
{
"queryshield": {
"command": "python",
"args": ["-m", "queryshield.mcp_server"],
"env": {
"QUERYSHIELD_API_KEY": "qs_...",
"QUERYSHIELD_BASE_URL": "https://api.queryshield.io"
}
}
}
```
暴露的工具：
- `query_database(database_alias, question, max_rows)` — 自然语言查询
- `query_database_sql(database_alias, sql, max_rows)` — 预构建的 SELECT
- `get_audit_log(limit)` — 调用 agent 的最近尝试记录
## 安全模型
| 威胁 | 防御手段 |
| ---------------------------------------------- | ------------------------------------------ |
| Agent 构造 `DROP TABLE` | sqlglot AST 拒绝非 SELECT 语句 |
| Agent 偷偷插入 `;` 和第二条语句 | 解析器拒绝 `len(statements) > 1` |
| Agent 使用 `pg_sleep`, `xp_cmdshell`, ... | AST 节点级别的函数拒绝列表 |
| Agent 读取其权限范围外的表 | RLS schema + table 白名单 |
| Agent 读取其他租户的数据行 | 通过 AST `.where()` 注入 `row_filters` |
| 连接字符串通过 stack trace 泄露 | Fernet 加密，绝不通过任何 API 返回 |
| 审计日志成为数据泄露向量 | 仅存储元数据 —— 绝不包含数据行 |
| `VAULT_KEY` 轮换 | 在新密钥下重新加密数据行（脚本驱动） |
`safety.py` 是最重要的模块。其中合并的每一项额外
检查，都应随附 `tests/test_safety.py` 中的测试。
## 定价
| 层级 | 月费 | 数据库 | 查询 / 月 | 备注 |
| ----------- | ------- | --------- | --------------- | ---------------------- |
| Starter | $500 | 3 | 1,000,000 | |
| Pro | $1,500 | 10 | 10,000,000 | 审计导出 |
| Enterprise | $3,500 | 不限 | 不限 | SSO, SIEM webhook |
目标为 `$32.5K MRR @ 15 customers (10 Pro + 5 Enterprise)`。
## 部署
该仓库已适配 Railway。`python -m queryshield.start` 是入口点
（通过 `os.getenv` 读取 `PORT`，因为 Railway 执行 start 命令时不带
shell）。从 Railway 市场配置 Postgres +（可选）Redis，
其余部分均通过环境变量设置。
```
railway up
```
`/health` 是存活检查。如果 control-plane
DB 无法连接，`/ready` 将返回 503。
## 测试
```
pip install pytest
python -m pytest tests/
```
42 个测试涵盖：
- AST 安全（24 个用例 —— 直接 DDL、注释、编码关键字、多条
语句、禁用函数、缺少 LIMIT）
- RLS 引擎（6 个用例 —— 白名单强制执行、WHERE 注入、与
现有谓词的连接）
- 针对 SQLite “客户 DB” 的代理端到端测试（5 个用例 —— 正常路径、
被拦截的 DML、RLS 行过滤、表白名单、缓存命中）
- 通过 FastAPI TestClient 进行 HTTP 集成（7 个用例 —— 完整的配置 →
查询流程、带 RLS 的受限 agent、授权失败）
标签：SQL代理, Streamlit, 安全规则引擎, 搜索引擎查询, 测试用例, 访问控制, 请求拦截, 逆向工具