From 9f5e01f721a2953b0d0d6599371b0e51dbc1630d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 21 Oct 2026 04:26:07 +0000 Subject: [PATCH] Fold AgentVault into QueryShield as the credential vault. Named secrets now live on the existing tenant and qs_ API key, encrypted with VAULT_KEY and gated by permission patterns. InjectShield stays a linked free funnel. The encryption test fixture no longer resembles a Stripe live key. Co-authored-by: brett-vikinetic --- .env.example | 15 +- CLAUDE.md | 25 +- README.md | 52 +- packages/queryshield-mcp/README.md | 6 +- .../queryshield-mcp/queryshield_mcp/server.py | 27 + queryshield/credential_api.py | 219 +++++++ queryshield/credential_crypto.py | 45 ++ queryshield/credentials.py | 543 ++++++++++++++++++ queryshield/main.py | 38 +- queryshield/mcp_http.py | 20 + queryshield/mcp_server.py | 35 ++ queryshield/models.py | 116 +++- queryshield/web.py | 156 ++++- server.json | 2 +- tests/test_credential_auth.py | 39 ++ tests/test_credentials.py | 376 ++++++++++++ tests/test_encryption.py | 63 ++ 17 files changed, 1734 insertions(+), 43 deletions(-) create mode 100644 queryshield/credential_api.py create mode 100644 queryshield/credential_crypto.py create mode 100644 queryshield/credentials.py create mode 100644 tests/test_credential_auth.py create mode 100644 tests/test_credentials.py create mode 100644 tests/test_encryption.py diff --git a/.env.example b/.env.example index 3e1fd45..81421f2 100644 --- a/.env.example +++ b/.env.example @@ -1,5 +1,5 @@ # --- Required: AI --- -ANTHROPIC_API_KEY=sk-ant-... +ANTHROPIC_API_KEY= # --- Required: internal control-plane DB (audit log + agents + databases + encrypted creds) --- # Railway provisions this — set DATABASE_URL automatically. @@ -7,7 +7,9 @@ DATABASE_URL=postgresql://user:pass@host:5432/queryshield # --- Required: vault encryption key (Fernet) --- # Generate once with: python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())" -# If this changes, all stored connection strings become unreadable. +# Encrypts database connection strings AND named credentials (the vault feature). +# There is one key. Losing it makes every stored secret unreadable. +# Named-credential routes fail closed when this is unset. Do not commit a key. VAULT_KEY= # --- Optional: cache --- @@ -15,14 +17,17 @@ VAULT_KEY= REDIS_URL= # --- Billing (Stripe) --- -STRIPE_SECRET_KEY=sk_live_... -STRIPE_WEBHOOK_SECRET=whsec_... +# One Stripe customer per workspace. Credential storage is included in the +# QueryShield prices below — do not create separate AgentVault products. +# InjectShield (the free funnel) keeps its own signup until you fold it. +STRIPE_SECRET_KEY= +STRIPE_WEBHOOK_SECRET= STRIPE_PRICE_STARTER=price_... STRIPE_PRICE_PRO=price_... STRIPE_PRICE_ENTERPRISE=price_... # --- Notifications --- -SENDGRID_API_KEY=SG.... +SENDGRID_API_KEY= SENDGRID_FROM=alerts@queryshield.io # Product-specific Discord webhook — separate channel from other agentic-build products. QUERY_SHIELD_DISCORD_WEBHOOK_URL= diff --git a/CLAUDE.md b/CLAUDE.md index 1f82e08..8e55b5f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,19 +1,25 @@ # QueryShield — Claude Code Context ## What this is -Secure SQL proxy between AI agents and customer databases. Agents POST to -`/v1/query` in NL or structured SQL. QueryShield validates → translates → -applies row-level security → executes → audits. Agents never see connection -strings. +Lead product of one agent security suite. Agents POST to `/v1/query` in NL +or structured SQL. QueryShield validates → translates → applies row-level +security → executes → audits. The same `qs_` API key and the same Stripe +customer also cover the named credential vault (`credentials.py`), ported +from AgentVault. InjectShield (https://github.com/bch1212/injectshield, +https://injectshield.dev) is the free prompt-injection funnel and is not +vendored here. Agents never see database connection strings. ## Architecture - `proxy.py` — engine. Every query flows through `execute_query()`. - `safety.py` — sqlglot AST validator. **Never use keyword filtering.** - `rls.py` — AST-level WHERE injection + schema/table whitelists. Never string concat. -- `vault.py` — Fernet-encrypted credential store, in our internal Postgres. - Replaces AWS Secrets Manager from the original spec (Brett's stack uses - Railway, no AWS). +- `vault.py` — Fernet-encrypted database connection strings, in our internal + Postgres. Replaces AWS Secrets Manager from the original spec (Brett's + stack uses Railway, no AWS). Never returned by any API. +- `credential_crypto.py` / `credentials.py` / `credential_api.py` — named + secret vault (AgentVault folded in). Same `VAULT_KEY`, same tenant, same + `qs_` key. Empty permission patterns fail closed. Audit metadata only. - `nl_to_sql.py` — Claude tool_use loop with `get_schema` + `validate_query` tools. Prompt caching on system + schema. - `schema.py` — DB introspection with two-tier cache (in-process + Redis). @@ -36,9 +42,10 @@ strings. in depth — even if safety + RLS were bypassed, the DB would refuse writes. ## Pricing tiers (`models.TIER_LIMITS`) -- Starter: $500/mo · 3 DBs · 1M queries/mo -- Pro: $1,500/mo · 10 DBs · 10M queries/mo · audit export +- Starter: $500/mo · 3 DBs · 25 credentials · 1M queries/mo +- Pro: $1,500/mo · 10 DBs · 100 credentials · 10M queries/mo · audit export - Enterprise: $3,500/mo · unlimited · SSO, SIEM webhook +- InjectShield is the free funnel, not a QueryShield Stripe price. ## MRR target $32.5K @ 15 customers (10 Pro + 5 Enterprise). First 3 customers = diff --git a/README.md b/README.md index 19b5e65..f8ca344 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [![tests](https://github.com/bch1212/queryshield/actions/workflows/test.yml/badge.svg)](https://github.com/bch1212/queryshield/actions/workflows/test.yml) -Secure SQL proxy between AI agents and enterprise databases. +Agent security suite. QueryShield is the lead product: a secure SQL proxy and a named credential vault, one `qs_` API key, one Stripe customer, and this docs site. Agents call a single endpoint in plain English (or structured SQL). QueryShield: @@ -17,6 +17,10 @@ Agents call a single endpoint in plain English (or structured SQL). QueryShield: Agents never see connection strings. +The credential vault stores other secrets (API tokens the agent is allowed to use). It is the AgentVault product, folded into this service. Permission patterns decide which names an agent can read. Database connection strings stay in a separate encrypted store and are not readable through the vault. + +**Free tier / funnel:** [InjectShield](https://injectshield.dev) scans untrusted text for prompt injection. It stays in its own repository ([bch1212/injectshield](https://github.com/bch1212/injectshield)) and its own API ([api.injectshield.dev](https://api.injectshield.dev)). Start there, then create one QueryShield workspace when an agent needs a database or a vault. Docs: [injectshield.dev/docs](https://injectshield.dev/docs). Suite page on this service: `/suite`. + --- ## Quickstart @@ -71,6 +75,31 @@ curl -X POST localhost:8000/v1/policies \ "row_filters": { "users": "tenant_id = 42" } }' +# 4b) Store a named credential and let one agent read it. +# The value is write-only on this call. Empty permissions grant nothing. +curl -X POST localhost:8000/v1/credentials \ + -H 'X-Admin-Key: qs_...' \ + -H 'Content-Type: application/json' \ + -d '{ + "name": "openai_key", + "provider": "openai", + "value": "YOUR_SECRET" + }' + +curl -X POST localhost:8000/v1/agents \ + -H 'X-Admin-Key: qs_...' \ + -H 'Content-Type: application/json' \ + -d '{ + "name": "writer", + "tenant_id": "", + "credential_permissions": ["openai_*"], + "budget_daily": 5 + }' + +# The agent uses the same qs_ key it uses for /v1/query. +curl -X POST localhost:8000/v1/vault/get/openai_key \ + -H 'X-API-Key: qs_...' + # 5) The agent queries. curl -X POST localhost:8000/v1/query \ -H 'X-API-Key: qs_...' \ @@ -131,8 +160,12 @@ Tools exposed: - `query_database(database_alias, question, max_rows)` — natural-language - `query_database_sql(database_alias, sql, max_rows)` — pre-built SELECT +- `list_credentials()` — names this agent may read +- `get_credential(name, cost)` — one permitted secret, same API key - `get_audit_log(limit)` — recent attempts for the calling agent +The PyPI package `queryshield-mcp` picks up the vault tools on its next release. Until then, `python -m queryshield.mcp_server` from this repo already exposes them. + --- ## Security model @@ -145,6 +178,8 @@ Tools exposed: | Agent reads tables outside its scope | RLS schema + table whitelist | | Agent reads other tenants' rows | `row_filters` injected via AST `.where()` | | Connection string leaks via stack traces | Fernet-encrypted, never returned in any API | +| Agent reads a secret it was not granted | fnmatch permissions; empty list fails closed | +| Vault audit becomes the secret | credential access log stores metadata only | | Audit log becomes the data exfil vector | only metadata is stored — never rows | | `VAULT_KEY` rotation | re-encrypt rows under new key (script-driven) | @@ -155,13 +190,14 @@ that lands there should ship with a test in `tests/test_safety.py`. ## Pricing -| Tier | Monthly | Databases | Queries / month | Notes | -| ----------- | ------- | --------- | --------------- | ---------------------- | -| Starter | $500 | 3 | 1,000,000 | | -| Pro | $1,500 | 10 | 10,000,000 | audit export | -| Enterprise | $3,500 | unlimited | unlimited | SSO, SIEM webhook | +| Tier | Monthly | Databases | Credentials | Queries / month | Notes | +| ------------ | ------- | --------- | ----------- | --------------- | ---------------------- | +| InjectShield | free | — | — | — | prompt-injection funnel, separate site | +| Starter | $500 | 3 | 25 | 1,000,000 | | +| Pro | $1,500 | 10 | 100 | 10,000,000 | audit export | +| Enterprise | $3,500 | unlimited | unlimited | unlimited | SSO, SIEM webhook | -Targets `$32.5K MRR @ 15 customers (10 Pro + 5 Enterprise)`. +QueryShield tiers are one Stripe customer. InjectShield sign-up stays on [injectshield.dev](https://injectshield.dev) until that billing is folded in. Targets `$32.5K MRR @ 15 customers (10 Pro + 5 Enterprise)`. --- @@ -188,7 +224,7 @@ pip install pytest python -m pytest tests/ ``` -42 tests cover: +100 tests cover the original safety, RLS, proxy, and HTTP suites plus the credential vault (encryption, `qs_` key hashing, permission checks, budgets, tenant isolation, and the column migration for existing databases). The original cases: - AST safety (24 cases — direct DDL, comments, encoded keywords, multiple statements, forbidden functions, missing LIMIT) diff --git a/packages/queryshield-mcp/README.md b/packages/queryshield-mcp/README.md index 4206aa1..d609abe 100644 --- a/packages/queryshield-mcp/README.md +++ b/packages/queryshield-mcp/README.md @@ -1,6 +1,6 @@ # queryshield-mcp -MCP client for [QueryShield](https://queryshield.dev) — a secure SQL proxy that sits between AI agents and customer databases. Agents call a single endpoint; QueryShield translates NL → SQL via Claude, validates at the AST level (no DDL/DML), enforces per-agent row-level security, and audit-logs every query. +MCP client for [QueryShield](https://queryshield.dev) — the agent security suite. One `QUERYSHIELD_API_KEY` reaches the SQL proxy and the named credential vault. [InjectShield](https://injectshield.dev) is the free prompt-injection funnel and stays a separate project. ## Install @@ -29,8 +29,12 @@ Self-hosting? Set `QUERYSHIELD_BASE_URL` to point at your deploy. - `query_database(database_alias, question, max_rows)` — natural-language query - `query_database_sql(database_alias, sql, max_rows)` — structured SELECT +- `list_credentials()` — credential names this agent may read +- `get_credential(name, cost=0)` — decrypt one permitted credential - `get_audit_log(limit)` — recent attempts for the calling agent +`list_credentials` and `get_credential` are in this source tree. They ship in the next PyPI release of `queryshield-mcp`. The Docker image used for MCP directory introspection still installs the published package until that release is out. + ## License MIT diff --git a/packages/queryshield-mcp/queryshield_mcp/server.py b/packages/queryshield-mcp/queryshield_mcp/server.py index 3585569..3d35fec 100644 --- a/packages/queryshield-mcp/queryshield_mcp/server.py +++ b/packages/queryshield-mcp/queryshield_mcp/server.py @@ -88,6 +88,33 @@ async def query_database_sql( return r.json() +@mcp.tool() +async def list_credentials() -> dict[str, Any]: + """List credential names this agent is allowed to read. Names only.""" + async with httpx.AsyncClient(timeout=30) as client: + r = await client.get( + f"{_base_url()}/v1/vault/credentials", + headers={"X-API-Key": _api_key()}, + ) + if r.status_code >= 400: + return {"error": r.text, "status": r.status_code} + return r.json() + + +@mcp.tool() +async def get_credential(name: str, cost: float = 0.0) -> dict[str, Any]: + """Read one named credential. Uses the same QUERYSHIELD_API_KEY as queries.""" + async with httpx.AsyncClient(timeout=30) as client: + r = await client.post( + f"{_base_url()}/v1/vault/get/{name}", + headers={"X-API-Key": _api_key()}, + params={"cost": cost}, + ) + if r.status_code >= 400: + return {"error": r.text, "status": r.status_code} + return r.json() + + @mcp.tool() async def get_audit_log(limit: int = 50) -> dict[str, Any]: """Recent audit entries for the calling agent.""" diff --git a/queryshield/credential_api.py b/queryshield/credential_api.py new file mode 100644 index 0000000..9d2b3df --- /dev/null +++ b/queryshield/credential_api.py @@ -0,0 +1,219 @@ +"""HTTP surface for the named credential vault. + +Admin routes use ``X-Admin-Key``. Agent routes use the same ``X-API-Key`` +as ``/v1/query`` — there is no second key prefix. Plaintext values are +returned only from the agent reveal route. +""" +from __future__ import annotations + +from typing import Optional + +from fastapi import APIRouter, Cookie, Depends, Form, HTTPException, Request, Response +from pydantic import BaseModel, Field + +from queryshield import auth as qs_auth +from queryshield.credentials import ( + CredentialError, + budget_summary, + create_credential, + delete_credential, + get_credential, + list_access_log, + list_available_names, + list_credentials, + reveal_credential, + update_agent_vault, + update_credential, +) +from queryshield.models import Agent +from queryshield.rate_limit import check as rl_check, client_ip + + +class CredentialCreateBody(BaseModel): + name: str = Field(..., min_length=1, max_length=255) + value: str = Field(..., min_length=1) + provider: Optional[str] = None + rotation_interval_hours: Optional[int] = None + metadata: dict = Field(default_factory=dict) + + +class CredentialUpdateBody(BaseModel): + name: Optional[str] = None + provider: Optional[str] = None + value: Optional[str] = None + rotation_interval_hours: Optional[int] = None + metadata: Optional[dict] = None + + +class AgentVaultUpdateBody(BaseModel): + name: Optional[str] = None + credential_permissions: Optional[list[str]] = None + budget_daily: Optional[float] = None + budget_monthly: Optional[float] = None + active: Optional[bool] = None + + +def _raise(exc: CredentialError) -> None: + raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc + + +def build_credential_router(require_admin, require_agent): # type: ignore[no-untyped-def] + router = APIRouter() + + @router.post("/v1/credentials", status_code=201) + async def create_named_credential( + body: CredentialCreateBody, admin: Agent = Depends(require_admin) + ): + try: + return create_credential( + admin.tenant_id, + body.name, + body.value, + provider=body.provider, + rotation_interval_hours=body.rotation_interval_hours, + metadata=body.metadata, + ) + except CredentialError as e: + _raise(e) + + @router.get("/v1/credentials") + async def list_named_credentials(admin: Agent = Depends(require_admin)): + return {"credentials": list_credentials(admin.tenant_id)} + + @router.get("/v1/credentials/access-log") + async def credential_access_log(limit: int = 100, admin: Agent = Depends(require_admin)): + return {"entries": list_access_log(admin.tenant_id, limit)} + + @router.get("/v1/credentials/{credential_id}") + async def get_named_credential(credential_id: str, admin: Agent = Depends(require_admin)): + try: + return get_credential(admin.tenant_id, credential_id) + except CredentialError as e: + _raise(e) + + @router.patch("/v1/credentials/{credential_id}") + async def patch_named_credential( + credential_id: str, + body: CredentialUpdateBody, + admin: Agent = Depends(require_admin), + ): + try: + return update_credential( + admin.tenant_id, + credential_id, + name=body.name, + provider=body.provider, + value=body.value, + rotation_interval_hours=body.rotation_interval_hours, + metadata=body.metadata, + fields_set=set(body.model_fields_set), + ) + except CredentialError as e: + _raise(e) + + @router.delete("/v1/credentials/{credential_id}", status_code=204) + async def remove_named_credential(credential_id: str, admin: Agent = Depends(require_admin)): + try: + delete_credential(admin.tenant_id, credential_id) + except CredentialError as e: + _raise(e) + return Response(status_code=204) + + @router.patch("/v1/agents/{agent_id}") + async def patch_agent_vault( + agent_id: str, + body: AgentVaultUpdateBody, + admin: Agent = Depends(require_admin), + ): + try: + return update_agent_vault( + admin.tenant_id, + agent_id, + name=body.name, + credential_permissions=body.credential_permissions, + budget_daily=body.budget_daily, + budget_monthly=body.budget_monthly, + active=body.active, + fields_set=set(body.model_fields_set), + ) + except CredentialError as e: + _raise(e) + + @router.get("/v1/agents/{agent_id}/budget") + async def get_agent_budget(agent_id: str, admin: Agent = Depends(require_admin)): + try: + return budget_summary(admin.tenant_id, agent_id) + except CredentialError as e: + _raise(e) + + @router.post("/v1/vault/get/{credential_name}") + async def vault_get( + credential_name: str, + request: Request, + cost: float = 0.0, + agent: Agent = Depends(require_agent), + ): + ip = client_ip(request) + allowed, retry = rl_check(ip, "vault_get", limit=120, window_sec=60) + if not allowed: + raise HTTPException( + status_code=429, + detail={"reason": "ip rate limit", "retry_after_sec": retry}, + headers={"Retry-After": str(retry)}, + ) + try: + return reveal_credential( + agent.id, + credential_name, + cost=cost, + ip_address=ip, + user_agent=request.headers.get("user-agent"), + ) + except CredentialError as e: + _raise(e) + + @router.get("/v1/vault/credentials") + async def vault_list(agent: Agent = Depends(require_agent)): + try: + return {"credentials": list_available_names(agent.id)} + except CredentialError as e: + _raise(e) + + def _tenant_from_session( + qs_session: str | None = Cookie(default=None, alias=qs_auth.SESSION_COOKIE), + ) -> str: + tenant_id = qs_auth.verify_session(qs_session) + if tenant_id is None: + raise HTTPException(status_code=401, detail="not authenticated") + return tenant_id + + @router.post("/dashboard/credentials") + async def dashboard_store_credential( + name: str = Form(...), + value: str = Form(...), + provider: str = Form(""), + tenant_id: str = Depends(_tenant_from_session), + ): + try: + created = create_credential( + tenant_id, + name, + value, + provider=provider or None, + ) + except CredentialError as e: + _raise(e) + return {"id": created["id"], "name": created["name"], "ok": True} + + @router.delete("/dashboard/credentials/{credential_id}") + async def dashboard_delete_credential( + credential_id: str, + tenant_id: str = Depends(_tenant_from_session), + ): + try: + delete_credential(tenant_id, credential_id) + except CredentialError as e: + _raise(e) + return {"id": credential_id, "deleted": True} + + return router diff --git a/queryshield/credential_crypto.py b/queryshield/credential_crypto.py new file mode 100644 index 0000000..f882948 --- /dev/null +++ b/queryshield/credential_crypto.py @@ -0,0 +1,45 @@ +"""Column-level Fernet encryption for named credentials. + +Uses ``VAULT_KEY``, the same key as the database-connection vault. There is +no baked-in fallback: a missing or invalid key refuses to encrypt or decrypt. +""" +from __future__ import annotations + +from cryptography.fernet import Fernet, InvalidToken + +from queryshield.config import get_settings + +_fernet: Fernet | None = None + + +def reset_fernet() -> None: + """Drop the cached Fernet instance. Tests call this after swapping keys.""" + global _fernet + _fernet = None + + +def _get_fernet() -> Fernet: + global _fernet + if _fernet is None: + key = get_settings().vault_key + if not key: + raise RuntimeError("VAULT_KEY is not set — cannot encrypt or decrypt credentials") + raw = key.encode("utf-8") if isinstance(key, str) else key + try: + _fernet = Fernet(raw) + except (TypeError, ValueError) as e: + raise RuntimeError("VAULT_KEY is not a valid Fernet key") from e + return _fernet + + +def encrypt_value(plaintext: str) -> str: + """Encrypt a credential value. Returns a urlsafe-base64 Fernet token.""" + return _get_fernet().encrypt(plaintext.encode("utf-8")).decode("utf-8") + + +def decrypt_value(ciphertext: str) -> str: + """Decrypt a credential value.""" + try: + return _get_fernet().decrypt(ciphertext.encode("utf-8")).decode("utf-8") + except InvalidToken as e: + raise ValueError("Failed to decrypt credential — key mismatch or corrupted data") from e diff --git a/queryshield/credentials.py b/queryshield/credentials.py new file mode 100644 index 0000000..8c33184 --- /dev/null +++ b/queryshield/credentials.py @@ -0,0 +1,543 @@ +"""Named credential vault for the QueryShield proxy. + +Ported from AgentVault and fitted onto the existing tenant, ``qs_`` agent +key, and ``VAULT_KEY``. Database connection strings stay in ``vault.py`` +and are never returned by this module. + +Permission patterns are fnmatch globs stored on the agent. An empty list +grants nothing. Audit rows hold metadata only — never the secret. +""" +from __future__ import annotations + +import fnmatch +import logging +import math +from datetime import datetime, timezone +from typing import Any, Optional + +from sqlalchemy import func, select + +from queryshield.credential_crypto import decrypt_value, encrypt_value +from queryshield.models import ( + Agent, + CredentialAccessLog, + CredentialBudgetUsage, + SessionLocal, + StoredCredential, + Tenant, + TIER_LIMITS, +) + +log = logging.getLogger("queryshield.credentials") + +TTL_SECONDS = 300 +_MAX_VALUE_CHARS = 100_000 +_MAX_DETAILS_CHARS = 8_000 + + +class CredentialError(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc) + + +def _permissions(raw: Any) -> list[str]: + if not raw: + return [] + if isinstance(raw, str): + return [] + return [str(p) for p in raw if isinstance(p, str) and p] + + +def credential_allowed(permissions: Any, credential_name: str) -> bool: + """True when ``credential_name`` matches a stored fnmatch pattern. + + Empty permissions fail closed. AgentVault treated an empty list as + "allow all"; that is unsafe once every existing QueryShield agent + would otherwise inherit every secret. + """ + patterns = _permissions(permissions) + if not patterns: + return False + return any(fnmatch.fnmatch(credential_name, pattern) for pattern in patterns) + + +def _public(row: StoredCredential) -> dict: + return { + "id": row.id, + "name": row.name, + "provider": row.provider, + "rotation_interval_hours": row.rotation_interval_hours, + "last_rotated": row.last_rotated.isoformat() if row.last_rotated else None, + "metadata": row.details or {}, + "created_at": row.created_at.isoformat() if row.created_at else None, + } + + +def _validate_name(name: str) -> str: + cleaned = (name or "").strip() + if not cleaned or len(cleaned) > 255: + raise CredentialError(400, "credential name must be 1–255 characters") + if "/" in cleaned or "\\" in cleaned or ".." in cleaned: + raise CredentialError(400, "credential name cannot contain path characters") + return cleaned + + +def _validate_value(value: str) -> str: + if not isinstance(value, str) or value == "": + raise CredentialError(400, "credential value is required") + if len(value) > _MAX_VALUE_CHARS: + raise CredentialError(400, "credential value is too large") + return value + + +def _validate_details(details: Optional[dict]) -> dict: + import json + + payload = details or {} + if not isinstance(payload, dict): + raise CredentialError(400, "metadata must be an object") + encoded = json.dumps(payload) + if len(encoded) > _MAX_DETAILS_CHARS: + raise CredentialError(400, "metadata is too large") + return payload + + +def _credential_limit(tenant_id: str, session) -> int: + tenant = session.get(Tenant, tenant_id) + tier = tenant.tier if tenant is not None else "starter" + limits = TIER_LIMITS.get(tier, TIER_LIMITS["starter"]) + return int(limits.get("credentials", 0)) + + +def create_credential( + tenant_id: str, + name: str, + value: str, + provider: Optional[str] = None, + rotation_interval_hours: Optional[int] = None, + metadata: Optional[dict] = None, +) -> dict: + name = _validate_name(name) + value = _validate_value(value) + details = _validate_details(metadata) + if provider is not None and len(provider) > 100: + raise CredentialError(400, "provider name is too long") + if rotation_interval_hours is not None and rotation_interval_hours < 1: + raise CredentialError(400, "rotation_interval_hours must be positive") + + try: + token = encrypt_value(value) + except RuntimeError as e: + raise CredentialError(500, str(e)) from e + + with SessionLocal() as session: + limit = _credential_limit(tenant_id, session) + count = session.execute( + select(func.count()).select_from(StoredCredential).where( + StoredCredential.tenant_id == tenant_id + ) + ).scalar_one() + if limit > 0 and count >= limit: + raise CredentialError( + 403, + f"Credential limit reached ({limit}). Upgrade your plan for more.", + ) + existing = session.execute( + select(StoredCredential).where( + StoredCredential.tenant_id == tenant_id, + StoredCredential.name == name, + ) + ).scalar_one_or_none() + if existing is not None: + raise CredentialError(409, f"Credential '{name}' already exists") + row = StoredCredential( + tenant_id=tenant_id, + name=name, + provider=provider, + encrypted_value=token, + rotation_interval_hours=rotation_interval_hours, + details=details, + ) + session.add(row) + session.commit() + log.info("credentials: stored %s/%s", tenant_id, name) + return _public(row) + + +def list_credentials(tenant_id: str) -> list[dict]: + with SessionLocal() as session: + rows = session.execute( + select(StoredCredential) + .where(StoredCredential.tenant_id == tenant_id) + .order_by(StoredCredential.created_at.desc()) + ).scalars().all() + return [_public(r) for r in rows] + + +def get_credential(tenant_id: str, credential_id: str) -> dict: + with SessionLocal() as session: + row = session.get(StoredCredential, credential_id) + if row is None or row.tenant_id != tenant_id: + raise CredentialError(404, "Credential not found") + return _public(row) + + +def update_credential( + tenant_id: str, + credential_id: str, + *, + name: Optional[str] = None, + provider: Optional[str] = None, + value: Optional[str] = None, + rotation_interval_hours: Optional[int] = None, + metadata: Optional[dict] = None, + fields_set: Optional[set[str]] = None, +) -> dict: + fields_set = fields_set or set() + with SessionLocal() as session: + row = session.get(StoredCredential, credential_id) + if row is None or row.tenant_id != tenant_id: + raise CredentialError(404, "Credential not found") + if "name" in fields_set and name is not None: + cleaned = _validate_name(name) + clash = session.execute( + select(StoredCredential).where( + StoredCredential.tenant_id == tenant_id, + StoredCredential.name == cleaned, + StoredCredential.id != row.id, + ) + ).scalar_one_or_none() + if clash is not None: + raise CredentialError(409, f"Credential '{cleaned}' already exists") + row.name = cleaned + if "provider" in fields_set: + if provider is not None and len(provider) > 100: + raise CredentialError(400, "provider name is too long") + row.provider = provider + if "value" in fields_set and value is not None: + try: + row.encrypted_value = encrypt_value(_validate_value(value)) + except RuntimeError as e: + raise CredentialError(500, str(e)) from e + row.last_rotated = _utcnow() + if "rotation_interval_hours" in fields_set: + if rotation_interval_hours is not None and rotation_interval_hours < 1: + raise CredentialError(400, "rotation_interval_hours must be positive") + row.rotation_interval_hours = rotation_interval_hours + if "metadata" in fields_set and metadata is not None: + row.details = _validate_details(metadata) + row.updated_at = _utcnow() + session.commit() + log.info("credentials: updated %s/%s", tenant_id, row.name) + return _public(row) + + +def delete_credential(tenant_id: str, credential_id: str) -> None: + with SessionLocal() as session: + row = session.get(StoredCredential, credential_id) + if row is None or row.tenant_id != tenant_id: + raise CredentialError(404, "Credential not found") + session.delete(row) + session.commit() + log.info("credentials: deleted %s/%s", tenant_id, row.name) + + +def _period_key(kind: str, now: Optional[datetime] = None) -> str: + now = now or _utcnow() + if kind == "daily": + return now.strftime("daily:%Y-%m-%d") + return now.strftime("monthly:%Y-%m") + + +def _usage(session, agent_id: str, kind: str, limit_amount: float) -> CredentialBudgetUsage: + key = _period_key(kind) + row = session.execute( + select(CredentialBudgetUsage).where( + CredentialBudgetUsage.agent_id == agent_id, + CredentialBudgetUsage.period_key == key, + ) + ).scalar_one_or_none() + if row is not None: + row.limit_amount = limit_amount + return row + row = CredentialBudgetUsage( + agent_id=agent_id, + period_key=key, + amount_used=0.0, + limit_amount=limit_amount, + ) + session.add(row) + session.flush() + return row + + +def _check_budget(session, agent: Agent, cost: float) -> tuple[bool, str]: + if cost <= 0: + return True, "ok" + if agent.budget_daily is not None: + daily = _usage(session, agent.id, "daily", agent.budget_daily) + if daily.amount_used + cost > daily.limit_amount: + return False, ( + f"Daily budget exceeded: ${daily.amount_used:.2f} used of " + f"${daily.limit_amount:.2f} limit" + ) + if agent.budget_monthly is not None: + monthly = _usage(session, agent.id, "monthly", agent.budget_monthly) + if monthly.amount_used + cost > monthly.limit_amount: + return False, ( + f"Monthly budget exceeded: ${monthly.amount_used:.2f} used of " + f"${monthly.limit_amount:.2f} limit" + ) + return True, "ok" + + +def _record_spend(session, agent: Agent, cost: float) -> None: + if cost <= 0: + return + if agent.budget_daily is not None: + daily = _usage(session, agent.id, "daily", agent.budget_daily) + daily.amount_used = round(daily.amount_used + cost, 6) + if agent.budget_monthly is not None: + monthly = _usage(session, agent.id, "monthly", agent.budget_monthly) + monthly.amount_used = round(monthly.amount_used + cost, 6) + + +def _log_access( + session, + *, + tenant_id: str, + agent_id: str, + credential_id: str, + success: bool, + error_message: Optional[str], + ip_address: Optional[str], + user_agent: Optional[str], +) -> None: + session.add( + CredentialAccessLog( + tenant_id=tenant_id, + agent_id=agent_id, + credential_id=credential_id, + action="read", + success=success, + error_message=error_message[:512] if error_message else None, + ip_address=ip_address, + user_agent=user_agent[:500] if user_agent else None, + ) + ) + + +def reveal_credential( + agent_id: str, + credential_name: str, + cost: float = 0.0, + ip_address: Optional[str] = None, + user_agent: Optional[str] = None, +) -> dict: + """Return a decrypted credential for an agent, or raise CredentialError. + + The plaintext is the response body only. It is not written to the audit + log or the application log. + """ + if not math.isfinite(cost) or cost < 0: + raise CredentialError(400, "cost must be a finite number >= 0") + name = _validate_name(credential_name) + + with SessionLocal() as session: + agent = session.get(Agent, agent_id) + if agent is None or not agent.active: + raise CredentialError(401, "invalid or inactive API key") + + row = session.execute( + select(StoredCredential).where( + StoredCredential.tenant_id == agent.tenant_id, + StoredCredential.name == name, + ) + ).scalar_one_or_none() + + if not credential_allowed(agent.credential_permissions, name): + if row is not None: + _log_access( + session, + tenant_id=agent.tenant_id, + agent_id=agent.id, + credential_id=row.id, + success=False, + error_message="Permission denied", + ip_address=ip_address, + user_agent=user_agent, + ) + session.commit() + raise CredentialError( + 403, f"Agent does not have permission to access '{name}'" + ) + + if row is None: + raise CredentialError(404, f"Credential '{name}' not found") + + allowed, reason = _check_budget(session, agent, cost) + if not allowed: + _log_access( + session, + tenant_id=agent.tenant_id, + agent_id=agent.id, + credential_id=row.id, + success=False, + error_message=reason, + ip_address=ip_address, + user_agent=user_agent, + ) + session.commit() + raise CredentialError(429, reason) + + try: + value = decrypt_value(row.encrypted_value) + except (ValueError, RuntimeError) as e: + _log_access( + session, + tenant_id=agent.tenant_id, + agent_id=agent.id, + credential_id=row.id, + success=False, + error_message="Failed to decrypt credential", + ip_address=ip_address, + user_agent=user_agent, + ) + session.commit() + log.warning("credentials: decrypt failed for %s/%s", agent.tenant_id, name) + raise CredentialError(500, "Failed to decrypt credential") from e + + _record_spend(session, agent, cost) + _log_access( + session, + tenant_id=agent.tenant_id, + agent_id=agent.id, + credential_id=row.id, + success=True, + error_message=None, + ip_address=ip_address, + user_agent=user_agent, + ) + session.commit() + return { + "credential_id": row.id, + "name": row.name, + "provider": row.provider, + "value": value, + "ttl_seconds": TTL_SECONDS, + } + + +def list_available_names(agent_id: str) -> list[str]: + with SessionLocal() as session: + agent = session.get(Agent, agent_id) + if agent is None or not agent.active: + raise CredentialError(401, "invalid or inactive API key") + rows = session.execute( + select(StoredCredential.name).where( + StoredCredential.tenant_id == agent.tenant_id + ) + ).all() + names = [r[0] for r in rows] + return [n for n in names if credential_allowed(agent.credential_permissions, n)] + + +def list_access_log(tenant_id: str, limit: int = 100) -> list[dict]: + limit = max(1, min(limit, 1000)) + with SessionLocal() as session: + rows = session.execute( + select(CredentialAccessLog) + .where(CredentialAccessLog.tenant_id == tenant_id) + .order_by(CredentialAccessLog.created_at.desc()) + .limit(limit) + ).scalars().all() + return [ + { + "id": r.id, + "agent_id": r.agent_id, + "credential_id": r.credential_id, + "action": r.action, + "ip_address": r.ip_address, + "success": r.success, + "error_message": r.error_message, + "created_at": r.created_at.isoformat() if r.created_at else None, + } + for r in rows + ] + + +def budget_summary(tenant_id: str, agent_id: str) -> dict: + with SessionLocal() as session: + agent = session.get(Agent, agent_id) + if agent is None or agent.tenant_id != tenant_id: + raise CredentialError(404, "agent not in this tenant") + summary: dict[str, Any] = { + "agent_id": agent.id, + "agent_name": agent.name, + "daily_limit": agent.budget_daily, + "monthly_limit": agent.budget_monthly, + "daily_used": 0.0, + "monthly_used": 0.0, + } + if agent.budget_daily is not None: + daily = _usage(session, agent.id, "daily", agent.budget_daily) + summary["daily_used"] = daily.amount_used + session.commit() + if agent.budget_monthly is not None: + monthly = _usage(session, agent.id, "monthly", agent.budget_monthly) + summary["monthly_used"] = monthly.amount_used + session.commit() + return summary + + +def update_agent_vault( + tenant_id: str, + agent_id: str, + *, + name: Optional[str] = None, + credential_permissions: Optional[list[str]] = None, + budget_daily: Optional[float] = None, + budget_monthly: Optional[float] = None, + active: Optional[bool] = None, + fields_set: Optional[set[str]] = None, +) -> dict: + fields_set = fields_set or set() + with SessionLocal() as session: + agent = session.get(Agent, agent_id) + if agent is None or agent.tenant_id != tenant_id: + raise CredentialError(404, "agent not in this tenant") + if "name" in fields_set and name is not None: + if not name.strip(): + raise CredentialError(400, "name is required") + agent.name = name.strip()[:255] + if "credential_permissions" in fields_set and credential_permissions is not None: + if not isinstance(credential_permissions, list) or any( + not isinstance(p, str) or not p or len(p) > 255 for p in credential_permissions + ): + raise CredentialError(400, "credential_permissions must be a list of patterns") + agent.credential_permissions = list(credential_permissions) + for field, value in (("budget_daily", budget_daily), ("budget_monthly", budget_monthly)): + if field in fields_set: + if value is not None and (not math.isfinite(value) or value < 0): + raise CredentialError(400, f"{field} must be >= 0") + setattr(agent, field, value) + if "active" in fields_set and active is not None: + if agent.name == "__admin__" and active is False: + raise CredentialError(400, "the admin agent cannot be deactivated") + agent.active = active + session.commit() + return { + "id": agent.id, + "name": agent.name, + "api_key_prefix": agent.api_key_prefix, + "credential_permissions": _permissions(agent.credential_permissions), + "budget_daily": agent.budget_daily, + "budget_monthly": agent.budget_monthly, + "active": agent.active, + } diff --git a/queryshield/main.py b/queryshield/main.py index 5aea45e..ff7f4d6 100644 --- a/queryshield/main.py +++ b/queryshield/main.py @@ -39,6 +39,7 @@ handle_webhook, increment_quota, ) +from queryshield.credential_api import build_credential_router from queryshield.config import get_settings from queryshield.models import ( Agent, @@ -125,9 +126,9 @@ async def lifespan(app: FastAPI): title="QueryShield", version=__version__, description=( - "Secure SQL proxy between AI agents and enterprise databases. " - "Validate intent, translate NL → SQL, enforce row-level security per " - "agent identity, and audit-log every query." + "QueryShield agent security suite. Secure SQL proxy plus a named " + "credential vault for AI agents. One qs_ API key, one Stripe customer. " + "InjectShield (https://injectshield.dev) is the free prompt-injection funnel." ), lifespan=lifespan, ) @@ -177,6 +178,9 @@ async def require_admin(x_admin_key: str = Header(..., alias="X-Admin-Key")) -> return agent +app.include_router(build_credential_router(require_admin, require_agent)) + + # --- Health ------------------------------------------------------------ @app.get("/health") @@ -321,13 +325,24 @@ async def create_tenant(request: Request, name: str): async def create_agent(reg: AgentRegistration, admin: Agent = Depends(require_admin)): if reg.tenant_id != admin.tenant_id: raise HTTPException(status_code=403, detail="tenant mismatch") + import math + raw, prefix, digest = generate_api_key() + for label, value in (("budget_daily", reg.budget_daily), ("budget_monthly", reg.budget_monthly)): + if value is not None and (not math.isfinite(value) or value < 0): + raise HTTPException(status_code=400, detail=f"{label} must be >= 0") + patterns = list(reg.credential_permissions or []) + if any(not isinstance(p, str) or not p.strip() for p in patterns): + raise HTTPException(status_code=400, detail="credential_permissions must be a list of patterns") with SessionLocal() as session: agent = Agent( tenant_id=reg.tenant_id, name=reg.name, api_key_hash=digest, api_key_prefix=prefix, + credential_permissions=patterns, + budget_daily=reg.budget_daily, + budget_monthly=reg.budget_monthly, ) session.add(agent) session.commit() @@ -434,6 +449,7 @@ async def billing_webhook(request: Request): LANDING_HTML, LOGIN_HTML, SIGNUP_RESULT_HTML, + SUITE_HTML, VERIFY_FAILED_HTML, render_guide, render_guides_index, @@ -470,6 +486,8 @@ def sitemap_xml() -> Response: '\n' " https://queryshield.dev/" "weekly1.0\n" + " https://queryshield.dev/suite" + "monthly0.8\n" " https://queryshield.dev/aeo/guides" "monthly0.7\n" + guide_urls @@ -480,6 +498,12 @@ def sitemap_xml() -> Response: return Response(content=xml, media_type="application/xml") +@app.get("/suite", response_class=HTMLResponse) +def suite_page() -> str: + """One docs page for the suite. InjectShield stays on its own site.""" + return SUITE_HTML + + @app.get("/aeo/guides", response_class=HTMLResponse) def aeo_guides_index() -> str: return render_guides_index() @@ -631,6 +655,9 @@ async def dashboard_data(tenant_id: str = Depends(_require_session)): "name": a.name, "key_prefix": a.api_key_prefix, "active": a.active, + "credential_permissions": a.credential_permissions or [], + "budget_daily": a.budget_daily, + "budget_monthly": a.budget_monthly, "created_at": a.created_at.isoformat() if a.created_at else None, } for a in agents @@ -638,7 +665,10 @@ async def dashboard_data(tenant_id: str = Depends(_require_session)): tier = tenant.tier or "starter" limits = TIER_LIMITS.get(tier, TIER_LIMITS["starter"]) + from queryshield.credentials import list_credentials as list_named_credentials + databases = list_databases(tenant_id) + named_credentials = list_named_credentials(tenant_id) logs = await _get_logs(tenant_id, limit=25) return { "tenant": { @@ -649,9 +679,11 @@ async def dashboard_data(tenant_id: str = Depends(_require_session)): "queries_used": tenant.queries_used_period, "queries_limit": limits["queries_per_month"], "databases_limit": limits["databases"], + "credentials_limit": limits.get("credentials", 0), }, "agents": agents_data, "databases": databases, + "credentials": named_credentials, "audit": [r.model_dump(mode="json") for r in logs], } diff --git a/queryshield/mcp_http.py b/queryshield/mcp_http.py index ad078fc..909d6a9 100644 --- a/queryshield/mcp_http.py +++ b/queryshield/mcp_http.py @@ -20,6 +20,7 @@ from queryshield.audit import get_agent_logs from queryshield.billing import check_quota, increment_quota +from queryshield.credentials import CredentialError, list_available_names, reveal_credential from queryshield.models import Agent, QueryRequest, SessionLocal, hash_api_key from queryshield.proxy import execute_query @@ -124,6 +125,25 @@ async def query_database_sql( increment_quota(agent.tenant_id, by=1) return result.model_dump(mode="json") + @mcp.tool() + async def list_credentials() -> dict[str, Any]: + """List credential names this agent is allowed to read.""" + agent = _agent_or_raise() + try: + names = list_available_names(agent.id) + except CredentialError as e: + return {"error": e.detail, "status": e.status_code} + return {"credentials": names} + + @mcp.tool() + async def get_credential(name: str, cost: float = 0.0) -> dict[str, Any]: + """Read one named credential this agent is permitted to use.""" + agent = _agent_or_raise() + try: + return reveal_credential(agent.id, name, cost=cost) + except CredentialError as e: + return {"error": e.detail, "status": e.status_code} + @mcp.tool() async def get_audit_log(limit: int = 50) -> dict[str, Any]: """Return recent audit entries for the calling agent.""" diff --git a/queryshield/mcp_server.py b/queryshield/mcp_server.py index db12b49..f75bd0f 100644 --- a/queryshield/mcp_server.py +++ b/queryshield/mcp_server.py @@ -92,6 +92,41 @@ async def query_database_sql( return r.json() +@mcp.tool() +async def list_credentials() -> dict[str, Any]: + """List credential names this agent is allowed to read. + + Names only — never values. Database connection strings are not included. + """ + async with httpx.AsyncClient(timeout=30) as client: + r = await client.get( + f"{_base_url()}/v1/vault/credentials", + headers={"X-API-Key": _api_key()}, + ) + if r.status_code >= 400: + return {"error": r.text, "status": r.status_code} + return r.json() + + +@mcp.tool() +async def get_credential(name: str, cost: float = 0.0) -> dict[str, Any]: + """Read one named credential the agent is permitted to use. + + ``cost`` is charged against the agent's daily and monthly vault budget + when those caps are set. The value is in the response only; discard it + when ``ttl_seconds`` elapses. QueryShield does not log the value. + """ + async with httpx.AsyncClient(timeout=30) as client: + r = await client.post( + f"{_base_url()}/v1/vault/get/{name}", + headers={"X-API-Key": _api_key()}, + params={"cost": cost}, + ) + if r.status_code >= 400: + return {"error": r.text, "status": r.status_code} + return r.json() + + @mcp.tool() async def get_audit_log(limit: int = 50) -> dict[str, Any]: """Recent audit entries for the calling agent. diff --git a/queryshield/models.py b/queryshield/models.py index 7cb01de..0332ea5 100644 --- a/queryshield/models.py +++ b/queryshield/models.py @@ -18,6 +18,7 @@ Boolean, Column, DateTime, + Float, Index, Integer, LargeBinary, @@ -43,9 +44,27 @@ class Tier(str, enum.Enum): TIER_LIMITS: Dict[str, Dict[str, int]] = { - "starter": {"databases": 3, "queries_per_month": 1_000_000, "price_cents": 50_000}, - "pro": {"databases": 10, "queries_per_month": 10_000_000, "price_cents": 150_000}, - "enterprise": {"databases": 10_000, "queries_per_month": 1_000_000_000, "price_cents": 350_000}, + # ``credentials`` is the named-secret vault (the folded AgentVault feature). + # Database connection strings are counted under ``databases``, not here. + # One Stripe customer on the QueryShield prices covers both. + "starter": { + "databases": 3, + "queries_per_month": 1_000_000, + "credentials": 25, + "price_cents": 50_000, + }, + "pro": { + "databases": 10, + "queries_per_month": 10_000_000, + "credentials": 100, + "price_cents": 150_000, + }, + "enterprise": { + "databases": 10_000, + "queries_per_month": 1_000_000_000, + "credentials": 10_000, + "price_cents": 350_000, + }, } @@ -93,6 +112,11 @@ class AgentRegistration(BaseModel): name: str tenant_id: str tier: Tier = Tier.starter + # fnmatch patterns for named credentials this agent may read. + # Empty means no credential access (fail closed). Use ["*"] for all. + credential_permissions: List[str] = Field(default_factory=list) + budget_daily: Optional[float] = None + budget_monthly: Optional[float] = None class AgentRegistrationResult(BaseModel): @@ -186,6 +210,10 @@ class Agent(Base): api_key_hash = Column(String(128), nullable=False, unique=True, index=True) api_key_prefix = Column(String(16), nullable=False) active = Column(Boolean, default=True, nullable=False) + # Named-credential scope. Empty list = no reads. Not a second API key. + credential_permissions = Column(JSON, nullable=False, default=list) + budget_daily = Column(Float, nullable=True) + budget_monthly = Column(Float, nullable=True) created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False) @@ -220,6 +248,62 @@ class RLSPolicyRow(Base): __table_args__ = (Index("ix_rls_agent_db", "agent_id", "database_alias", unique=True),) +class StoredCredential(Base): + """Tenant-scoped named secret. ``encrypted_value`` is a Fernet token. + + This is the AgentVault credential, stored on the QueryShield tenant. + Database connection strings do not live here — see ``DatabaseRow``. + """ + + __tablename__ = "stored_credentials" + + id = Column(String(64), primary_key=True, default=_new_id) + tenant_id = Column(String(64), nullable=False, index=True) + name = Column(String(255), nullable=False) + provider = Column(String(100), nullable=True) + encrypted_value = Column(Text, nullable=False) + rotation_interval_hours = Column(Integer, nullable=True) + last_rotated = Column(DateTime(timezone=True), nullable=True) + details = Column(JSON, nullable=False, default=dict) + created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False) + updated_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False) + + __table_args__ = (Index("ix_cred_tenant_name", "tenant_id", "name", unique=True),) + + +class CredentialAccessLog(Base): + """Append-only metadata for named-credential reads. Never the secret.""" + + __tablename__ = "credential_access_logs" + + id = Column(String(64), primary_key=True, default=_new_id) + tenant_id = Column(String(64), nullable=False, index=True) + agent_id = Column(String(64), nullable=False, index=True) + credential_id = Column(String(64), nullable=False, index=True) + action = Column(String(50), nullable=False) + ip_address = Column(String(45), nullable=True) + user_agent = Column(String(500), nullable=True) + success = Column(Boolean, default=True, nullable=False) + error_message = Column(String(512), nullable=True) + created_at = Column(DateTime(timezone=True), default=_utcnow, nullable=False, index=True) + + +class CredentialBudgetUsage(Base): + """Per-agent spend against a daily or monthly cap. ``period_key`` is UTC.""" + + __tablename__ = "credential_budget_usage" + + id = Column(String(64), primary_key=True, default=_new_id) + agent_id = Column(String(64), nullable=False, index=True) + period_key = Column(String(32), nullable=False) + amount_used = Column(Float, default=0.0, nullable=False) + limit_amount = Column(Float, nullable=False) + + __table_args__ = ( + Index("ix_budget_agent_period", "agent_id", "period_key", unique=True), + ) + + class AuditLog(Base): """Append-only. Never store rows — only metadata.""" @@ -263,14 +347,30 @@ def init_db() -> None: def _apply_column_migrations() -> None: from sqlalchemy import text - # (table, column, type) — Postgres-flavored. Each statement uses - # ADD COLUMN IF NOT EXISTS so the boot is safe to retry. - migrations = [ - ("tenants", "owner_email", "VARCHAR(320)"), - ] + # (table, column, type). Postgres uses ADD COLUMN IF NOT EXISTS. + # SQLite checks PRAGMA table_info and skips columns that already exist. dialect = ENGINE.dialect.name # "postgresql" | "sqlite" | ... if dialect not in {"postgresql", "sqlite"}: return + migrations = [ + ("tenants", "owner_email", "VARCHAR(320)"), + ] + if dialect == "postgresql": + migrations.extend( + [ + ("agents", "credential_permissions", "JSON"), + ("agents", "budget_daily", "DOUBLE PRECISION"), + ("agents", "budget_monthly", "DOUBLE PRECISION"), + ] + ) + else: + migrations.extend( + [ + ("agents", "credential_permissions", "TEXT"), + ("agents", "budget_daily", "REAL"), + ("agents", "budget_monthly", "REAL"), + ] + ) with ENGINE.begin() as conn: for table, column, coltype in migrations: if dialect == "postgresql": diff --git a/queryshield/web.py b/queryshield/web.py index d453653..9020652 100644 --- a/queryshield/web.py +++ b/queryshield/web.py @@ -158,6 +158,7 @@ QueryShield Log in + Suite API Guides GitHub @@ -171,6 +172,7 @@

SELECT-only AST validator Per-agent RLS + Credential vault Append-only audit MCP-native

@@ -185,7 +187,7 @@

-

Free tier: 3 databases · 1M queries/month · no credit card.

+

One qs_ API key for the SQL proxy and the credential vault. Creating a workspace does not charge a card; Starter limits apply until you subscribe, and that subscription is the only Stripe customer. The free funnel is InjectShield.

What is a secure SQL proxy for AI agents?

@@ -231,17 +233,98 @@

Browse all guides →

+

The agent security suite

+

QueryShield is the lead product: one brand, one qs_ API key, one Stripe customer, and this site. Two capabilities ship here.

+
    +
  • SQL proxy. Natural language or SELECT in, AST checks, row-level security, and an audit log. Agents never see database connection strings.
    • +
  • Credential vault. Named secrets (the former AgentVault) encrypted with the same VAULT_KEY. An agent reads only the names its permission patterns allow, through POST /v1/vault/get/{name}.
    • +
+

The free tier is InjectShield, a prompt-injection scanner for untrusted text. It stays a separate service and funnels teams here when they need a database proxy or a credential vault. Source and docs: github.com/bch1212/injectshield · injectshield.dev/docs · API api.injectshield.dev.

+

Suite overview

+

Pricing

- + - - - + + + +
TierMonthlyDatabasesQueries/month
TierMonthlyDatabasesCredentialsQueries/month
Starter$50031M
Pro$1,5001010M (audit export)
Enterprise$3,500unlimitedSSO + SIEM webhook
InjectShieldFree——prompt-injection scans
Starter$5003251M
Pro$1,5001010010M (audit export)
Enterprise$3,500unlimitedunlimitedSSO + SIEM webhook
+

QueryShield tiers share one Stripe customer. InjectShield sign-up is separate and free: injectshield.dev.

-

Docs: /docs · Health: /health · github.com/bch1212/queryshield

+

Suite: /suite · API: /docs · Health: /health · github.com/bch1212/queryshield

+ + + +""" +) + + +SUITE_HTML = ( + """ + + + + QueryShield suite — SQL proxy, credential vault, InjectShield + + + + + +
+ + +

One suite

+

QueryShield is the lead product. Agents use one qs_ API key. Billing is one Stripe customer on the QueryShield prices. This site is the docs.

+ +

SQL proxy

+

Register a database with the admin key. The connection string is Fernet-encrypted and never returned by the API. Agents call POST /v1/query in natural language or as a single SELECT. The AST validator, row-level security, and the append-only query audit run before any rows leave the database.

+ +

Credential vault

+

Named secrets — the AgentVault product, now a feature of this proxy. Store them with the admin key. Grant an agent fnmatch patterns such as stripe_*. The agent reads a value with its existing API key. An empty pattern list grants nothing.

+ 
# Store (admin). The response has metadata only.
+curl -X POST $BASE/v1/credentials \\
+  -H "X-Admin-Key: qs_..." \\
+  -H "Content-Type: application/json" \\
+  -d '{"name":"openai_key","provider":"openai","value":"YOUR_SECRET"}'
+
+# Let one agent read it. Same key family as /v1/query.
+curl -X POST $BASE/v1/agents \\
+  -H "X-Admin-Key: qs_..." \\
+  -H "Content-Type: application/json" \\
+  -d '{"name":"writer","tenant_id":"TENANT","credential_permissions":["openai_key"],"budget_daily":5}'
+
+# Reveal (agent). ttl_seconds is advisory — discard the value after use.
+curl -X POST $BASE/v1/vault/get/openai_key \\
+  -H "X-API-Key: qs_..."
+

Database connection strings are a different store. /v1/vault/get cannot read them. Access logs record the agent, credential id, and outcome, never the secret. Encryption key: VAULT_KEY only. If it is unset, credential routes fail closed.

+

MCP tools list_credentials and get_credential use QUERYSHIELD_API_KEY. OpenAPI for every route is at /docs.

+ +

Free tier: InjectShield

+

InjectShield scans untrusted text for prompt injection before it reaches an agent. It is the free funnel into this suite. Its code stays in its own repository.

+ +

A team starts with a free InjectShield key, then creates one QueryShield workspace when agents need a database or a vault. QueryShield does not issue a second key prefix and does not open a second Stripe customer for the vault.

+ +

What stayed behind

+

AgentVault's separate user directory, Clerk login, avk_ keys, and its own Stripe products (Free / Pro $49 / Business $149 / Enterprise $499) are not part of this service. Budgets and credential caps ride on the QueryShield tiers you already sell.

@@ -481,6 +564,28 @@ def replace(self, key, value="", *a, **kw): # type: ignore[override] +

Credential vault

+

Named secrets for agents (API keys, tokens). Values are encrypted at rest and are not shown again after you save them. Grant access with credential permission patterns on the API — an agent with an empty list cannot read any secret. Database connection strings stay in the section above and are never returned to agents.

+
+ + + +
NameProviderCreated
+
+ + Store a credential +
+ + + + + + + +
+
+
+
+

Recent activity

@@ -530,6 +635,17 @@ def replace(self, key, value="", *a, **kw): # type: ignore[override] `).join("") || ``; + const creds = document.querySelector("#cred-table tbody"); + const credRows = d.credentials || []; + creds.innerHTML = credRows.map(c => ` + + + + + + + `).join("") || ``; + const dbs = document.querySelector("#db-table tbody"); dbs.innerHTML = d.databases.map(x => ` @@ -597,6 +713,28 @@ def replace(self, key, value="", *a, **kw): # type: ignore[override] load(); } +async function storeCredential() { + const name = document.getElementById("cred-name").value.trim(); + const provider = document.getElementById("cred-provider").value.trim(); + const value = document.getElementById("cred-value").value; + if (!name || !value) { alert("Name and value are required."); return; } + const fd = new FormData(); + fd.append("name", name); fd.append("provider", provider); fd.append("value", value); + const r = await fetch("/dashboard/credentials", { method: "POST", body: fd, credentials: "same-origin" }); + if (!r.ok) { document.getElementById("cred-result").innerHTML = `

${esc(await r.text())}

`; return; } + document.getElementById("cred-result").innerHTML = `

Stored. The value is not shown again.

`; + document.getElementById("cred-name").value = ""; + document.getElementById("cred-provider").value = ""; + document.getElementById("cred-value").value = ""; + load(); +} + +async function deleteCredential(id, name) { + if (!confirm(`Remove credential "${name}"? Agents will no longer be able to read it.`)) return; + await fetch("/dashboard/credentials/" + encodeURIComponent(id), { method: "DELETE", credentials: "same-origin" }); + load(); +} + async function deleteDb(alias) { if (!confirm(`Remove database "${alias}"? Stored credentials will be deleted.`)) return; await fetch("/dashboard/databases/" + encodeURIComponent(alias), { method: "DELETE", credentials: "same-origin" }); @@ -636,8 +774,10 @@ def replace(self, key, value="", *a, **kw): # type: ignore[override] 'natural language in, SELECT-only validated SQL out, per-agent row-level ' 'security, and an append-only audit log. Your agents never see connection ' 'strings.

' - '

Get an API key — free tier ' - 'Read the API docs

' + '

Get a QueryShield API key ' + 'Suite overview

' + '

Free prompt-injection scans: ' + 'InjectShield.

' '' ) diff --git a/server.json b/server.json index 8739ef6..81831b7 100644 --- a/server.json +++ b/server.json @@ -1,7 +1,7 @@ { "$schema": "https://static.modelcontextprotocol.io/schemas/2025-09-29/server.schema.json", "name": "io.github.bch1212/queryshield", - "description": "Secure SQL proxy for AI agents — NL→SQL, AST safety, per-agent RLS, audit log.", + "description": "QueryShield agent security suite — SQL proxy, named credential vault, AST safety, per-agent RLS, audit log. InjectShield is the free prompt-injection funnel.", "version": "1.0.1", "repository": { "url": "https://github.com/bch1212/queryshield", diff --git a/tests/test_credential_auth.py b/tests/test_credential_auth.py new file mode 100644 index 0000000..b6920a6 --- /dev/null +++ b/tests/test_credential_auth.py @@ -0,0 +1,39 @@ +"""API key checks adapted from AgentVault's tests/test_auth.py. + +AgentVault issued a second prefix (avk_). The suite has one key: qs_, +generated and hashed by queryshield.models. +""" +from __future__ import annotations + +from queryshield.models import generate_api_key, hash_api_key + + +class TestApiKeyGeneration: + def test_key_has_prefix(self): + key, prefix, _digest = generate_api_key() + assert key.startswith("qs_") + assert prefix == key[:8] + assert not key.startswith("avk_") + + def test_keys_are_unique(self): + keys = {generate_api_key()[0] for _ in range(100)} + assert len(keys) == 100 + + def test_key_length(self): + key, _prefix, _digest = generate_api_key() + assert len(key) > 40 + + def test_hash_is_deterministic(self): + key, _prefix, digest = generate_api_key() + assert hash_api_key(key) == digest + assert hash_api_key(key) == hash_api_key(key) + + def test_hash_differs_for_different_keys(self): + k1, _, _ = generate_api_key() + k2, _, _ = generate_api_key() + assert hash_api_key(k1) != hash_api_key(k2) + + def test_prefix_extraction(self): + _key, prefix, _digest = generate_api_key() + assert prefix.startswith("qs_") + assert len(prefix) == 8 diff --git a/tests/test_credentials.py b/tests/test_credentials.py new file mode 100644 index 0000000..81a8038 --- /dev/null +++ b/tests/test_credentials.py @@ -0,0 +1,376 @@ +"""Named credential vault — AgentVault flows on the QueryShield API. + +Adapted from AgentVault's tests/test_api.py. Users, Clerk, avk_ keys, and +a second Stripe catalog are not part of this suite, so those cases are +expressed with a tenant admin key and a qs_ agent key. +""" +from __future__ import annotations + +import pytest +from fastapi.testclient import TestClient + +from queryshield.main import app + +# Obviously fake. Not a provider key. +_SECRET = "dummy_credential_value" +_OTHER = "dummy_other_credential_value" + + +@pytest.fixture() +def client(control_db): + return TestClient(app) + + +def _tenant(client: TestClient) -> tuple[str, str]: + r = client.post("/v1/tenants", params={"name": "Acme"}) + assert r.status_code == 200, r.text + body = r.json() + return body["api_key"], body["tenant_id"] + + +def _store(client: TestClient, admin_key: str, name: str, value: str, provider: str = "example") -> dict: + r = client.post( + "/v1/credentials", + headers={"X-Admin-Key": admin_key}, + json={"name": name, "provider": provider, "value": value}, + ) + assert r.status_code == 201, r.text + return r.json() + + +def _agent( + client: TestClient, + admin_key: str, + tenant_id: str, + *, + name: str = "worker", + permissions: list[str] | None = None, + budget_daily: float | None = None, + budget_monthly: float | None = None, +) -> str: + payload: dict = { + "name": name, + "tenant_id": tenant_id, + "credential_permissions": ["*"] if permissions is None else permissions, + } + if budget_daily is not None: + payload["budget_daily"] = budget_daily + if budget_monthly is not None: + payload["budget_monthly"] = budget_monthly + r = client.post("/v1/agents", headers={"X-Admin-Key": admin_key}, json=payload) + assert r.status_code == 200, r.text + return r.json()["api_key"] + + +def test_suite_page_links_injectshield(client) -> None: + r = client.get("/suite") + assert r.status_code == 200 + assert "InjectShield" in r.text + assert "https://github.com/bch1212/injectshield" in r.text + assert "https://injectshield.dev" in r.text + assert "credential" in r.text.lower() + + +def test_landing_mentions_suite(client) -> None: + r = client.get("/") + assert r.status_code == 200 + assert "InjectShield" in r.text + assert "Credential vault" in r.text + + +def test_create_credential_hides_value(client) -> None: + admin, _tenant_id = _tenant(client) + body = _store(client, admin, "openai_key", _SECRET, provider="openai") + assert body["name"] == "openai_key" + assert body["provider"] == "openai" + assert "value" not in body + assert _SECRET not in str(body) + + +def test_duplicate_credential_name(client) -> None: + admin, _tenant_id = _tenant(client) + _store(client, admin, "dup_key", _SECRET) + r = client.post( + "/v1/credentials", + headers={"X-Admin-Key": admin}, + json={"name": "dup_key", "value": _OTHER}, + ) + assert r.status_code == 409 + + +def test_credential_plan_limit(client, monkeypatch) -> None: + from queryshield.models import TIER_LIMITS + + monkeypatch.setitem(TIER_LIMITS["starter"], "credentials", 2) + admin, _tenant_id = _tenant(client) + _store(client, admin, "key_0", _SECRET) + _store(client, admin, "key_1", _OTHER) + r = client.post( + "/v1/credentials", + headers={"X-Admin-Key": admin}, + json={"name": "key_2", "value": "dummy_third"}, + ) + assert r.status_code == 403 + assert "limit" in r.text.lower() + + +def test_list_and_update_credential(client) -> None: + admin, _tenant_id = _tenant(client) + created = _store(client, admin, "update_me", _SECRET) + listed = client.get("/v1/credentials", headers={"X-Admin-Key": admin}) + assert listed.status_code == 200 + assert len(listed.json()["credentials"]) == 1 + assert _SECRET not in listed.text + + renamed = client.patch( + f"/v1/credentials/{created['id']}", + headers={"X-Admin-Key": admin}, + json={"name": "updated_name"}, + ) + assert renamed.status_code == 200, renamed.text + assert renamed.json()["name"] == "updated_name" + assert "value" not in renamed.json() + + +def test_get_credential_via_vault(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "stripe_key", _SECRET, provider="example") + _store(client, admin, "openai_key", _OTHER, provider="example") + _store(client, admin, "twilio_sid", "dummy_twilio_value", provider="example") + key = _agent( + client, + admin, + tenant_id, + permissions=["stripe_*", "openai_key"], + ) + r = client.post("/v1/vault/get/stripe_key", headers={"X-API-Key": key}) + assert r.status_code == 200, r.text + body = r.json() + assert body["value"] == _SECRET + assert body["name"] == "stripe_key" + assert body["ttl_seconds"] == 300 + + +def test_permission_denied(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "twilio_sid", "dummy_twilio_value") + key = _agent(client, admin, tenant_id, permissions=["stripe_*", "openai_key"]) + r = client.post("/v1/vault/get/twilio_sid", headers={"X-API-Key": key}) + assert r.status_code == 403 + + +def test_empty_permissions_fail_closed(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "test_key", _SECRET) + key = _agent(client, admin, tenant_id, permissions=[]) + r = client.post("/v1/vault/get/test_key", headers={"X-API-Key": key}) + assert r.status_code == 403 + + +def test_credential_not_found(client) -> None: + admin, tenant_id = _tenant(client) + key = _agent(client, admin, tenant_id, permissions=["stripe_*"]) + r = client.post("/v1/vault/get/stripe_nonexistent", headers={"X-API-Key": key}) + assert r.status_code == 404 + + +def test_invalid_agent_key(client) -> None: + r = client.post( + "/v1/vault/get/stripe_key", + headers={"X-API-Key": "qs_not_a_real_agent_key"}, + ) + assert r.status_code == 401 + + +def test_list_available_credentials(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "stripe_key", _SECRET) + _store(client, admin, "openai_key", _OTHER) + _store(client, admin, "twilio_sid", "dummy_twilio_value") + key = _agent(client, admin, tenant_id, permissions=["stripe_*", "openai_key"]) + r = client.get("/v1/vault/credentials", headers={"X-API-Key": key}) + assert r.status_code == 200, r.text + names = r.json()["credentials"] + assert "stripe_key" in names + assert "openai_key" in names + assert "twilio_sid" not in names + + +def test_database_connection_is_not_a_vault_credential(client, sample_customer_db) -> None: + admin, tenant_id = _tenant(client) + r = client.post( + "/v1/databases", + headers={"X-Admin-Key": admin}, + json={ + "alias": "prod", + "db_type": "sqlite", + "connection_string": sample_customer_db, + }, + ) + assert r.status_code == 200, r.text + key = _agent(client, admin, tenant_id, permissions=["*"]) + revealed = client.post("/v1/vault/get/prod", headers={"X-API-Key": key}) + assert revealed.status_code == 404 + assert sample_customer_db not in revealed.text + listed = client.get("/v1/databases", headers={"X-Admin-Key": admin}) + assert sample_customer_db not in listed.text + named = client.get("/v1/credentials", headers={"X-Admin-Key": admin}) + assert named.json()["credentials"] == [] + + +def test_budget_enforcement_blocks_overspend(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "test_key", _SECRET) + key = _agent( + client, + admin, + tenant_id, + permissions=["test_key"], + budget_daily=1.00, + budget_monthly=10.00, + ) + ok = client.post( + "/v1/vault/get/test_key", + headers={"X-API-Key": key}, + params={"cost": 0.80}, + ) + assert ok.status_code == 200, ok.text + blocked = client.post( + "/v1/vault/get/test_key", + headers={"X-API-Key": key}, + params={"cost": 0.30}, + ) + assert blocked.status_code == 429 + assert "budget" in blocked.json()["detail"].lower() + assert _SECRET not in blocked.text + + +def test_zero_cost_always_allowed(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "test_key", _SECRET) + key = _agent( + client, + admin, + tenant_id, + permissions=["*"], + budget_daily=1.00, + ) + for _ in range(5): + r = client.post( + "/v1/vault/get/test_key", + headers={"X-API-Key": key}, + params={"cost": 0.0}, + ) + assert r.status_code == 200, r.text + + +def test_access_log_omits_secret(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "openai_key", _SECRET) + key = _agent(client, admin, tenant_id, permissions=["openai_key"]) + assert client.post("/v1/vault/get/openai_key", headers={"X-API-Key": key}).status_code == 200 + log = client.get("/v1/credentials/access-log", headers={"X-Admin-Key": admin}) + assert log.status_code == 200, log.text + assert log.json()["entries"][0]["success"] is True + assert _SECRET not in log.text + + +def test_vault_with_cost_tracking(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "stripe_key", _SECRET) + key = _agent(client, admin, tenant_id, permissions=["stripe_*"]) + r = client.post( + "/v1/vault/get/stripe_key", + headers={"X-API-Key": key}, + params={"cost": 0.50}, + ) + assert r.status_code == 200, r.text + + +def test_cross_tenant_isolation(client) -> None: + admin_a, tenant_a = _tenant(client) + admin_b, tenant_b = _tenant(client) + created = _store(client, admin_a, "openai_key", _SECRET) + missing = client.get( + f"/v1/credentials/{created['id']}", + headers={"X-Admin-Key": admin_b}, + ) + assert missing.status_code == 404 + key_b = _agent(client, admin_b, tenant_b, permissions=["*"]) + revealed = client.post("/v1/vault/get/openai_key", headers={"X-API-Key": key_b}) + assert revealed.status_code == 404 + assert _SECRET not in revealed.text + key_a = _agent(client, admin_a, tenant_a, permissions=["openai_key"]) + ok = client.post("/v1/vault/get/openai_key", headers={"X-API-Key": key_a}) + assert ok.status_code == 200 + assert ok.json()["value"] == _SECRET + + +def test_dashboard_store_does_not_echo_value(client) -> None: + from queryshield import auth as qs_auth + from queryshield.models import SessionLocal, Tenant + + client.post("/signup", data={"email": "vault@example.com", "workspace": "Vault"}) + with SessionLocal() as session: + tenant = session.query(Tenant).filter_by(owner_email="vault@example.com").one() + tenant_id = tenant.id + token = qs_auth.issue_magic_link("vault@example.com", tenant_id) + assert client.get(f"/auth/verify?token={token}", follow_redirects=False).status_code == 303 + + stored = client.post( + "/dashboard/credentials", + data={"name": "openai_key", "provider": "openai", "value": _SECRET}, + ) + assert stored.status_code == 200, stored.text + assert stored.json()["name"] == "openai_key" + assert _SECRET not in stored.text + + page = client.get("/dashboard") + assert page.status_code == 200 + assert "Credential vault" in page.text + assert _SECRET not in page.text + + data = client.get("/dashboard/data") + assert data.status_code == 200 + assert data.json()["credentials"][0]["name"] == "openai_key" + assert _SECRET not in data.text + + +def test_negative_cost_rejected(client) -> None: + admin, tenant_id = _tenant(client) + _store(client, admin, "test_key", _SECRET) + key = _agent(client, admin, tenant_id, permissions=["*"]) + r = client.post( + "/v1/vault/get/test_key", + headers={"X-API-Key": key}, + params={"cost": -1}, + ) + assert r.status_code == 400 + assert _SECRET not in r.text + + +def test_existing_database_gains_vault_columns(control_db) -> None: + """create_all does not add columns to a table that already exists.""" + from sqlalchemy import text + + from queryshield.models import ENGINE, init_db + + with ENGINE.begin() as conn: + conn.execute(text("ALTER TABLE agents DROP COLUMN credential_permissions")) + conn.execute(text("ALTER TABLE agents DROP COLUMN budget_daily")) + conn.execute(text("ALTER TABLE agents DROP COLUMN budget_monthly")) + init_db() + with ENGINE.connect() as conn: + cols = {row[1] for row in conn.execute(text("PRAGMA table_info(agents)"))} + assert "credential_permissions" in cols + assert "budget_daily" in cols + assert "budget_monthly" in cols + + +def test_credential_allowed_patterns() -> None: + from queryshield.credentials import credential_allowed + + assert credential_allowed([], "openai_key") is False + assert credential_allowed(["openai_*"], "openai_key") is True + assert credential_allowed(["stripe_*"], "openai_key") is False + assert credential_allowed(["*"], "anything") is True diff --git a/tests/test_encryption.py b/tests/test_encryption.py new file mode 100644 index 0000000..957eb04 --- /dev/null +++ b/tests/test_encryption.py @@ -0,0 +1,63 @@ +"""Column-level encryption for named credentials. + +Ported from AgentVault's tests/test_encryption.py. The old fixture used a +string shaped like a Stripe live key; secret scanners flagged it. The value +below is an obvious dummy. +""" +from __future__ import annotations + +import pytest + +from queryshield.credential_crypto import decrypt_value, encrypt_value, reset_fernet + + +class TestEncryption: + def test_encrypt_decrypt_roundtrip(self): + """Encrypted value decrypts back to original.""" + reset_fernet() + secret = "dummy_not_a_real_stripe_key" + encrypted = encrypt_value(secret) + + assert encrypted != secret + assert decrypt_value(encrypted) == secret + + def test_different_ciphertexts_per_call(self): + """Fernet produces unique ciphertexts (nonce-based).""" + reset_fernet() + secret = "my_api_key_value" + a = encrypt_value(secret) + b = encrypt_value(secret) + assert a != b + + def test_decrypt_bad_data_raises(self): + """Decrypting garbage raises ValueError.""" + reset_fernet() + with pytest.raises(ValueError, match="Failed to decrypt"): + decrypt_value("not_valid_fernet_token") + + def test_empty_string_roundtrip(self): + """Can encrypt and decrypt empty string.""" + reset_fernet() + encrypted = encrypt_value("") + assert decrypt_value(encrypted) == "" + + def test_long_value_roundtrip(self): + """Can handle long values (e.g., JWT tokens).""" + reset_fernet() + long_val = "x" * 10000 + encrypted = encrypt_value(long_val) + assert decrypt_value(encrypted) == long_val + + def test_missing_vault_key_fails_closed(self, monkeypatch): + """No baked-in key. Encryption refuses to run without VAULT_KEY.""" + import queryshield.credential_crypto as crypto + + reset_fernet() + + class _Unset: + vault_key = None + + monkeypatch.setattr(crypto, "get_settings", lambda: _Unset()) + with pytest.raises(RuntimeError, match="VAULT_KEY"): + encrypt_value("dummy_not_a_real_stripe_key") + reset_fernet()
No agents yet.
${esc(c.name)}${esc(c.provider || "")}${c.created_at ? c.created_at.slice(0,10) : ""}
No named credentials yet.