From eec827a5bdaefbf1dcfd8a4230f6f86c71bf895e Mon Sep 17 00:00:00 2001 From: Deep Kumar Singh Kushwah Date: Mon, 8 Jun 2026 09:28:56 +0530 Subject: [PATCH] =?UTF-8?q?feat(crypto):=20zero-downtime=20key=20rotation?= =?UTF-8?q?=20=E2=80=94=20RotatingKeyManager=20(Step=206.7d)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close Step 6.7. RotatingKeyManager (rag-core, a crypto-free KeyManager decorator) makes rotating a tenant's KEK seamless: encrypt always uses the current key; decrypt tries the current key then each non-expired retired key (RetiredKey(key_manager, expires_at)). The try-all decrypt is safe — AES-GCM authenticates the DEK, so a wrong KEK can't yield a valid key (only the KEK that wrapped a blob decrypts it). rewrap(ctx, ct) is the background-migration primitive (decrypt with whatever key still works → re-encrypt under the current key); once all blobs are migrated a retired key is dropped. An expired retired key is skipped, so old un-rewrapped data is sealed (EncryptionError) with no impact on current-key data. Injectable clock for deterministic expiry tests. Composes with every provider (local + all four cloud KMS) since it orchestrates only the KeyManager SPI. ragctl kms --rotate demos the full flow. Scope: the rotation mechanism; config-driven per-tenant rotation + the storage re-encryption job land with the EncryptingStorage ingest wiring (tiered storage). Co-Authored-By: Claude Opus 4.7 --- TRACKER.md | 24 +-- docs/README.md | 2 +- docs/adr/ADR-0039-byok-envelope-encryption.md | 11 +- docs/architecture/byok.md | 8 +- docs/reference/encryption.md | 31 ++++ .../core/src/rag_core/rotating_key_manager.py | 103 +++++++++++++ packages/ragctl/src/ragctl/main.py | 41 +++++- packages/ragctl/tests/test_kms.py | 10 ++ tests/kms/test_rotation.py | 137 ++++++++++++++++++ 9 files changed, 350 insertions(+), 17 deletions(-) create mode 100644 packages/core/src/rag_core/rotating_key_manager.py create mode 100644 tests/kms/test_rotation.py diff --git a/TRACKER.md b/TRACKER.md index 2cbe592..19b37c4 100644 --- a/TRACKER.md +++ b/TRACKER.md @@ -14,15 +14,13 @@ | | | |---|---| | **Last updated** | 2026-06-08 | -| **Current phase** | Phase 6 — Governance & Tenancy (**6 / 10 steps**) | -| **Overall** | **70 / 84 steps** — Phases 0–5 complete | -| **Next action** | **Step 6.7d — Key rotation**: zero-downtime KEK rotation — old keys retained decrypt-only until expiry, background re-encryption. Closes Step 6.7. 6.7a (library + local) + 6.7b (config/factory + AWS KMS) + 6.7c (GCP/Azure/Vault providers) shipped. | +| **Current phase** | Phase 6 — Governance & Tenancy (**7 / 10 steps**) | +| **Overall** | **71 / 84 steps** — Phases 0–5 complete | +| **Next action** | **Step 6.8 — SSO / SCIM**: OIDC + SAML IdP federation; SCIM 2.0 user provisioning; per-tenant IdP config. New ground (the gateway's `Auth` SPI + `NoopAuth` header-identity is the seam this builds on). | **Recently shipped** -- **6.7c** 🚧 GCP/Azure/Vault KMS providers — `GcpKmsKeyManager` (google-cloud-kms), `AzureKeyVaultKeyManager` (`wrap_key`/`unwrap_key`), `VaultKeyManager` (Vault Transit) — each subclasses `EnvelopeKeyManager` behind a `[kms-*]` extra (lazy SDK, injectable client, fake-client unit-tested); factory + enum extended; completes all four KMS providers — [#157](https://github.com/officialCodeWork/AgentContextOS/pull/157) -- **6.7b** 🚧 Cloud KMS (AWS) + config/factory/wiring — `cfg.kms` (provider `noop`/`local`/`aws`) + `tenants[].kms_key_id`; `build_key_manager_from_config` factory + `app.state.key_manager` gateway seam; **`AwsKmsKeyManager`** (aioboto3, per-tenant CMK, KMS-error → sealing, fake-client unit-tested) — [#156](https://github.com/officialCodeWork/AgentContextOS/pull/156) -- **6.7a** 🚧 BYOK envelope encryption library — `KeyManager` SPI + `NoopKeyManager` + `EncryptingStorage` (rag-core); `EnvelopeKeyManager` (AES-256-GCM DEK + `tenant_id` AAD) + `LocalKeyManager` (in-process per-tenant KEK) (rag-backends); per-tenant isolation + **sealing** (`KeyUnavailableError`) + tamper detection; `ragctl kms` — [#155](https://github.com/officialCodeWork/AgentContextOS/pull/155) +- **6.7** ✅ BYOK envelope encryption — `KeyManager` SPI + `EncryptingStorage` + `EnvelopeKeyManager` (AES-256-GCM DEK + tenant AAD); `LocalKeyManager` + four cloud KMS providers (`Aws`/`Gcp`/`AzureKeyVault`/`Vault`, behind `[kms-*]` extras); `cfg.kms` + per-tenant key + factory; per-tenant isolation + sealing + tamper-evidence; zero-downtime rotation (`RotatingKeyManager` + `rewrap`); `ragctl kms` — [#155](https://github.com/officialCodeWork/AgentContextOS/pull/155)–[#158](https://github.com/officialCodeWork/AgentContextOS/pull/158) - **6.6** ✅ Immutable audit log — **6.6a** shared `AuditWriter`/store on `app.state` + `GET /v1/audit` (tenant-scoped, `chain_verified`) + `GET /v1/audit/verify` (whole-log) + `cfg.audit.enabled` ([#153](https://github.com/officialCodeWork/AgentContextOS/pull/153)); **6.6b** `AuditExporter` self-verifying WORM bundle (SHA-256 `content_hash` + HMAC, offline `verify()`), `POST /v1/audit/export` (tenant-scoped) + `ragctl audit`, `cfg.audit.export_secret` ([#154](https://github.com/officialCodeWork/AgentContextOS/pull/154)) - **6.5** ✅ PII egress policies — `PiiPolicyEngine` (`rag-pii`) answers `egress_text` over the context (`list[Chunk]`) + agent answer (`str`) the gateway already passes, applying the per-tenant `pii_policy` (allow / redact / mask / block); reuses the Step 1.7 detector + rewriters; opt-in `cfg.pii.enabled`; emits `pii.egress_blocked` — [#152](https://github.com/officialCodeWork/AgentContextOS/pull/152) - **6.4** ✅ ACL egress verifier — `AclEgressVerifier` re-checks returned `ChunkRef`s against the principal's labels at the gateway router boundary (defense-in-depth, same overlap semantics, independent of the PDP); `acl.egress_violation` event; red-team zero-violation-rate gate — [#151](https://github.com/officialCodeWork/AgentContextOS/pull/151) @@ -63,9 +61,9 @@ | 3 | Gateway & Agent Runtime | 11 | **11** | 0 | | 4 | Reliability | 6 | **6** | 0 | | 5 | Eval & Observability | 7 | **7** | 0 | -| 6 | Governance & Tenancy | 10 | **6** | 4 | +| 6 | Governance & Tenancy | 10 | **7** | 3 | | 7 | Pilot, Harden, GA | 10 | 0 | 10 | -| **Total** | | **84** | **70** | **14** | +| **Total** | | **84** | **71** | **13** | --- @@ -658,7 +656,7 @@ | 6.4 | ACL egress verifier | ✅ | [#151](https://github.com/officialCodeWork/AgentContextOS/pull/151) — `AclEgressVerifier` re-checks returned chunks at the gateway router boundary (defense-in-depth above the 6.3 push-down); `acl.egress_violation`; zero-violation-rate red-team gate | | 6.5 | PII policies | ✅ | [#152](https://github.com/officialCodeWork/AgentContextOS/pull/152) — `PiiPolicyEngine` egress_text decorator (allow / redact / mask / block per tenant) over answer + context; reuses Step 1.7 detector; `pii.egress_blocked` | | 6.6 | Immutable audit log | ✅ | **6.6a** [#153](https://github.com/officialCodeWork/AgentContextOS/pull/153) — read API `GET /v1/audit` + `GET /v1/audit/verify` + shared store + `cfg.audit`. **6.6b** [#154](https://github.com/officialCodeWork/AgentContextOS/pull/154) — `AuditExporter` signed WORM bundle + `POST /v1/audit/export` + `ragctl audit` | -| 6.7 | BYOK (Bring Your Own Key) | 🚧 | **6.7a** ✅ [#155](https://github.com/officialCodeWork/AgentContextOS/pull/155) — `KeyManager` SPI + envelope encryption (`LocalKeyManager`) + `EncryptingStorage` + sealing/isolation/tamper. **6.7b** ✅ [#156](https://github.com/officialCodeWork/AgentContextOS/pull/156) — `cfg.kms` + per-tenant key + `build_key_manager_from_config` factory + `AwsKmsKeyManager`. **6.7c** ✅ [#157](https://github.com/officialCodeWork/AgentContextOS/pull/157) — `GcpKmsKeyManager` / `AzureKeyVaultKeyManager` / `VaultKeyManager` behind `[kms-*]` extras. **6.7d** ⏳ — key rotation | +| 6.7 | BYOK (Bring Your Own Key) | ✅ | **6.7a** ✅ [#155](https://github.com/officialCodeWork/AgentContextOS/pull/155) — `KeyManager` SPI + envelope encryption (`LocalKeyManager`) + `EncryptingStorage` + sealing/isolation/tamper. **6.7b** ✅ [#156](https://github.com/officialCodeWork/AgentContextOS/pull/156) — `cfg.kms` + per-tenant key + `build_key_manager_from_config` factory + `AwsKmsKeyManager`. **6.7c** ✅ [#157](https://github.com/officialCodeWork/AgentContextOS/pull/157) — `GcpKmsKeyManager` / `AzureKeyVaultKeyManager` / `VaultKeyManager` behind `[kms-*]` extras. **6.7d** ✅ [#158](https://github.com/officialCodeWork/AgentContextOS/pull/158) — `RotatingKeyManager` zero-downtime rotation + `rewrap` | | 6.8 | SSO / SCIM | ⏳ | OIDC + SAML IdP federation; SCIM 2.0 user provisioning; per-tenant IdP config | | 6.9 | Air-gapped install bundle | ⏳ | Signed tarball with all images + Helm chart; offline bootstrap; cosign verification | | 6.10 | Compliance posture | ⏳ | SOC 2 Type II control mapping; GDPR data-residency config; data-retention policies | @@ -754,6 +752,13 @@ New ground — the only prior crypto was HMAC signing. The V1 plan calls for en - Each is **behind a `[kms-gcp]` / `[kms-azure]` / `[kms-vault]` extra** (the SDKs are lazy-imported, so the modules import without the SDK and selecting a provider without its extra raises a clear `ImportError`); an **injectable `client` / `client_factory` seam** makes every provider fully unit-testable with a fake KMS — no cloud creds / network. Connection uses each SDK's standard credential discovery (GCP ADC, Azure `DefaultAzureCredential`, Vault `VAULT_ADDR`/`VAULT_TOKEN`); per-tenant key resolution + sealing (`KeyUnavailableError`) are uniform with AWS - `KmsProvider` enum + `build_key_manager_from_config` factory + `KmsConfig.vault_mount` extended. mypy overrides added for the three SDK module trees (absent at lint time). **Scope:** providers only — key rotation is **6.7d**; wiring `EncryptingStorage` into the ingest path stays deferred (tiered storage). ~18 new tests (each provider over a fake client: round-trip / cross-tenant block / sealing / SDK-error mapping; factory cloud-provider-requires-extra contract). `KmsConfig` → `rag.schema` regenerated; all gates green (ruff, mypy --strict 309 files, RAG001, policy-coverage, log-schema). [ADR-0039](docs/adr/ADR-0039-byok-envelope-encryption.md), [reference/encryption.md](docs/reference/encryption.md), [architecture/byok.md](docs/architecture/byok.md) +#### 6.7d — Zero-downtime key rotation ✅ [#158](https://github.com/officialCodeWork/AgentContextOS/pull/158) + +- Closes Step 6.7. New **`RotatingKeyManager`** (`rag-core`, a **crypto-free** `KeyManager` decorator) makes rotating a tenant's KEK seamless: **encrypt** always uses the **current** key; **decrypt** tries the current key then each **non-expired** retired key (`RetiredKey(key_manager, expires_at)`). The try-all decrypt is **safe, not a guess** — AES-GCM authenticates the DEK, so a wrong KEK can't yield a valid key (only the KEK that wrapped a blob decrypts it) +- **`rewrap(ctx, ct)`** is the background-migration primitive (decrypt with whatever key still works → re-encrypt under the current key); once all blobs are migrated a retired key is dropped. An **expired** retired key is skipped, so old un-rewrapped data is **sealed** (`EncryptionError`) — "retain decrypt-only until expiry" — with no impact on current-key data. Injectable `clock` for deterministic expiry tests +- Composes with every provider (local + all four cloud KMS) since it orchestrates only the `KeyManager` SPI. `ragctl kms --rotate` demos the full flow (old + new decrypt, rewrap, expiry seals old data) +- **Scope:** the rotation *mechanism* (satisfies the planning rotation acceptance — rotate → old + new both decrypt; expired key seals old data). **Deferred:** config-driven per-tenant multi-generation rotation + the storage-side background re-encryption job land with the `EncryptingStorage` ingest wiring (tiered storage). ~11 new tests (rotate → old+new decrypt; rewrap migration; expired-key sealing; tamper; no-key sealing; multi-generation; clock-controlled expiry; health) + `ragctl kms --rotate`. No dist drift; all gates green (ruff, mypy --strict 310 files, RAG001, policy-coverage, log-schema). [ADR-0039](docs/adr/ADR-0039-byok-envelope-encryption.md), [reference/encryption.md](docs/reference/encryption.md), [architecture/byok.md](docs/architecture/byok.md) + --- ## Phase 7 — Pilot, Harden, GA (Weeks 34–40) ⏳ @@ -915,6 +920,7 @@ Complete log of every PR. Routine Dependabot bumps are grouped; everything else | [#155](https://github.com/officialCodeWork/AgentContextOS/pull/155) | 2026-06-08 | feat(crypto): BYOK envelope encryption library + local KMS (Step 6.7a) | | [#156](https://github.com/officialCodeWork/AgentContextOS/pull/156) | 2026-06-08 | feat(crypto): cfg.kms + key-manager factory + AWS KMS provider (Step 6.7b) | | [#157](https://github.com/officialCodeWork/AgentContextOS/pull/157) | 2026-06-08 | feat(crypto): GCP / Azure / Vault KMS providers (Step 6.7c) | +| [#158](https://github.com/officialCodeWork/AgentContextOS/pull/158) | 2026-06-08 | feat(crypto): zero-downtime key rotation — RotatingKeyManager (Step 6.7d) | | #78–#80, #116–#118 | Open | Dependabot bumps — awaiting merge | | #81 | Closed | Dependabot bump — superseded | diff --git a/docs/README.md b/docs/README.md index 9bbcdd1..25bcc28 100644 --- a/docs/README.md +++ b/docs/README.md @@ -107,7 +107,7 @@ | [admin-ui.md](reference/admin-ui.md) | Admin console (Step 3.10) — Next.js 14 operator GUI (`apps/admin-ui`); 9 pages (dashboard, corpora, connectors, glossary, webhooks, audit, API keys, tenants, config), live-vs-seed hybrid + `NEXT_PUBLIC_GATEWAY_URL`, header identity, running it, internals (shell/primitives/data layer), extension points | | [tenancy.md](reference/tenancy.md) | Logical multi-tenancy (Step 6.1) — per-tenant `rag.yaml` config (`namespace` / `acl_labels` / `pii_policy` / `quota`); `TenantResolver.resolve(id) → TenantSettings`; `RequestContext.namespace`; `GET /v1/status/tenant`; `ragctl tenant list` / `resolve`; config table + scope/boundaries (6.2/6.3/6.5) + extension points; physical tenancy (6.2), ACL push-down (6.3) + egress verifier (6.4 — `cfg.acl.verify_egress`) sections | | [audit.md](reference/audit.md) | Audit log (Step 6.6) — `AuditEvent` / `AuditStore` (append / events / verify_chain) / `NoopAuditStore` SHA-256 hash chain / `AuditWriter` (+ `.store`); read API `GET /v1/audit` (tenant-scoped, `chain_verified`) + `GET /v1/audit/verify` (whole-log); WORM signed export (6.6b) — `AuditExporter` (content_hash + HMAC), `POST /v1/audit/export`, offline `verify()`, `ragctl audit`; `cfg.audit.enabled` / `export_secret`; durable-store extension points | -| [encryption.md](reference/encryption.md) | BYOK envelope encryption (Step 6.7) — `KeyManager` SPI + `NoopKeyManager`; `EnvelopeKeyManager` (AES-256-GCM DEK + `tenant_id` AAD) + `LocalKeyManager` + cloud providers **`AwsKmsKeyManager`** / **`GcpKmsKeyManager`** / **`AzureKeyVaultKeyManager`** / **`VaultKeyManager`** (behind `[kms-*]` extras); `EncryptingStorage` decorator; `EncryptionError` / `KeyUnavailableError` (sealing); provider table; `cfg.kms` + `tenants[].kms_key_id` + `build_key_manager_from_config` factory; `ragctl kms`; rotation deferred (6.7d) | +| [encryption.md](reference/encryption.md) | BYOK envelope encryption (Step 6.7) — `KeyManager` SPI + `NoopKeyManager`; `EnvelopeKeyManager` (AES-256-GCM DEK + `tenant_id` AAD) + `LocalKeyManager` + cloud providers **`AwsKmsKeyManager`** / **`GcpKmsKeyManager`** / **`AzureKeyVaultKeyManager`** / **`VaultKeyManager`** (behind `[kms-*]` extras); `EncryptingStorage` decorator; `EncryptionError` / `KeyUnavailableError` (sealing); provider table; `cfg.kms` + `tenants[].kms_key_id` + `build_key_manager_from_config` factory; **`RotatingKeyManager`** + `RetiredKey` + `rewrap` (6.7d zero-downtime rotation); `ragctl kms [--rotate]` | | [webhooks.md](reference/webhooks.md) | Outbound webhooks (Step 3.9) — event catalogue (`ingest.completed` / `audit.policy_violation` / `drift.detected` / `eval.regression`), event envelope, HMAC signing + `verify()`, at-least-once delivery, `/v1/webhooks/subscriptions` CRUD + test, `rag.yaml` block, `ragctl webhooks demo`, internals + extension points | | [integrations.md](reference/integrations.md) | Framework adapters (Step 3.8) — `agentcontextos.integrations.*` for LangChain / LlamaIndex / Haystack / DSPy / LangGraph / CrewAI / AutoGen / Semantic Kernel; per-framework extras, shared config + chunk metadata, usage per framework, internals + extension points | | [status-api.md](reference/status-api.md) | Status & Metrics API (Step 3.11) — `/v1/status/health` / `metrics` / `logs` (+ SSE `logs/stream`), `WS /v1/status/ws`, `/v1/connectors/status`; metric catalogue + request-timing middleware, the `MetricsCollector` / `LogTail` read-side, CORS + query-param identity for browser streams, extension points | diff --git a/docs/adr/ADR-0039-byok-envelope-encryption.md b/docs/adr/ADR-0039-byok-envelope-encryption.md index bbeaf6a..9053603 100644 --- a/docs/adr/ADR-0039-byok-envelope-encryption.md +++ b/docs/adr/ADR-0039-byok-envelope-encryption.md @@ -81,9 +81,16 @@ where the real provider impls live. rag-core stays crypto-free (SPI + decorator (`AzureKeyVaultKeyManager`, `wrap_key`/`unwrap_key`), **HashiCorp Vault Transit** (`VaultKeyManager`) — each behind a `[kms-*]` extra (lazy SDK, injectable client, fake-client unit-tested), completing all four KMS providers. +- **6.7d** added zero-downtime key rotation: `RotatingKeyManager` (a crypto-free + `KeyManager` decorator) encrypts with the current key and decrypts via the + current key then each **non-expired** retired key (`RetiredKey`), with `rewrap` + for background migration. The try-all decrypt is safe — AES-GCM authenticates + the DEK, so a wrong KEK never yields a valid key — and an expired retired key + seals old, un-rewrapped data. This satisfies the rotation acceptance criteria. - Still deferred: wiring `EncryptingStorage` into the ingest/storage path (blocked - on tiered-storage plumbing — chunk content is inline today) and **6.7d** - zero-downtime key rotation (old keys retained decrypt-only until expiry). + on tiered-storage plumbing — chunk content is inline today) and config-driven + rotation (the `RotatingKeyManager` mechanism ships; per-tenant multi-generation + rotation config + the background re-encryption job land with that wiring). ## See also - [architecture/byok.md](../architecture/byok.md) — envelope design, AAD, sealing, slicing diff --git a/docs/architecture/byok.md b/docs/architecture/byok.md index c7f499a..d34f30e 100644 --- a/docs/architecture/byok.md +++ b/docs/architecture/byok.md @@ -87,8 +87,12 @@ asserts its reads fail while another tenant's succeed. **Azure Key Vault** (`AzureKeyVaultKeyManager`, via `wrap_key`/`unwrap_key`), **HashiCorp Vault Transit** (`VaultKeyManager`), each behind a `[kms-*]` extra (lazy SDK, injectable client). Completes all four KMS providers. -- **6.7d** — zero-downtime key rotation (old keys retained decrypt-only until - expiry; background re-encryption). +- **6.7d** — zero-downtime key rotation: `RotatingKeyManager` (a crypto-free + `KeyManager` decorator) encrypts with the current key, decrypts via current + + non-expired retired keys, and `rewrap`s old blobs to the current key. The + try-all decrypt is safe (AES-GCM authenticates the DEK); an expired retired key + seals old data. Config-driven rotation + the storage re-encryption job land with + the `EncryptingStorage` ingest wiring. ## Boundary note diff --git a/docs/reference/encryption.md b/docs/reference/encryption.md index 6be1fbe..d7b1dc6 100644 --- a/docs/reference/encryption.md +++ b/docs/reference/encryption.md @@ -94,6 +94,37 @@ missing key → `KeyUnavailableError` (sealing). A tenant with neither its own `kms_key_id` nor `default_key_id` is sealed. The SDKs are lazy-imported, so selecting a provider without its extra raises a clear `ImportError`. +### Key rotation (Step 6.7d) + +`RotatingKeyManager` (`rag_core.rotating_key_manager`) wraps a **current** +`KeyManager` plus **retired** keys for zero-downtime KEK rotation: + +```python +from datetime import UTC, datetime, timedelta +from rag_core.rotating_key_manager import RetiredKey, RotatingKeyManager + +km = RotatingKeyManager( + current=new_key_manager, + retired=[RetiredKey(old_key_manager, expires_at=datetime.now(UTC) + timedelta(days=30))], +) +new_ct = await km.encrypt(ctx, data) # always the current key +plain = await km.decrypt(ctx, any_ct) # current, then non-expired retired keys +migrated = await km.rewrap(ctx, old_ct) # re-encrypt an old blob under current +``` + +- **encrypt** uses the current key; **decrypt** tries the current key then each + non-expired retired key. The try-all is *safe* (AES-GCM authenticates the DEK — + a wrong KEK can't yield a valid key, so it never returns wrong plaintext). +- An **expired** retired key is skipped, so old data not yet re-wrapped is sealed + (raises `EncryptionError`) — "retain decrypt-only until expiry". +- **rewrap** is the background-migration primitive: decrypt with whatever key + still works, re-encrypt under the current key; once all blobs are migrated the + retired key can be dropped. `ragctl kms --rotate` demos the full flow. + +It is crypto-free (it orchestrates the `KeyManager` SPI), so it composes with the +local + every cloud provider. Config-driven rotation + the storage-side +re-encryption job land with the `EncryptingStorage` ingest wiring. + ## Internals - **Envelope wire format:** `b"RAGK" | version | u32(len(wrapped_dek)) | diff --git a/packages/core/src/rag_core/rotating_key_manager.py b/packages/core/src/rag_core/rotating_key_manager.py new file mode 100644 index 0000000..1c24db6 --- /dev/null +++ b/packages/core/src/rag_core/rotating_key_manager.py @@ -0,0 +1,103 @@ +"""RotatingKeyManager — zero-downtime KEK rotation (Step 6.7d). + +A :class:`~rag_core.spi.key_manager.KeyManager` decorator that makes rotating a +tenant's key-encryption key (KEK) seamless: new data is encrypted with the +**current** key, while data wrapped under **retired** keys keeps decrypting until +those keys expire. A background job re-wraps old blobs to the current key via +:meth:`rewrap`, after which a retired key can be dropped. + +* **encrypt** → always the current key. +* **decrypt** → try the current key, then each **non-expired** retired key. This + try-all is *safe*, not a guess: AES-GCM authenticates the data key, so a wrong + KEK fails to unwrap the DEK (it never yields a wrong-but-valid key) — only the + KEK that wrapped a blob can decrypt it. An expired retired key is skipped, so + old data not yet re-wrapped becomes **sealed** (``KeyUnavailableError``) — the + intended "retain decrypt-only until expiry" behaviour. +* **rewrap** → decrypt with whatever key still works, re-encrypt with the current + key (the background migration primitive). + +The decorator is **crypto-free** — it only orchestrates the wrapped key managers +(local or any cloud KMS) — so it composes with every provider. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime + +from rag_core.errors import EncryptionError, KeyUnavailableError +from rag_core.spi.key_manager import KeyManager +from rag_core.types import RequestContext + +__all__ = ["RetiredKey", "RotatingKeyManager"] + + +@dataclass(frozen=True) +class RetiredKey: + """A previously-current key manager, kept decrypt-only until ``expires_at``. + + ``expires_at`` is timezone-aware; ``None`` means "no expiry" (retained until + removed from the rotation). Once expired, the key is no longer tried on + decrypt, so any blob still wrapped under it becomes unreadable (sealed). + """ + + key_manager: KeyManager + expires_at: datetime | None = None + + +def _now_utc() -> datetime: + return datetime.now(UTC) + + +class RotatingKeyManager(KeyManager): + """Decorate a current ``KeyManager`` with retired keys for zero-downtime rotation.""" + + def __init__( + self, + *, + current: KeyManager, + retired: Sequence[RetiredKey] = (), + clock: Callable[[], datetime] | None = None, + ) -> None: + self._current = current + self._retired = list(retired) + self._clock = clock or _now_utc + + @property + def current(self) -> KeyManager: + return self._current + + async def encrypt(self, ctx: RequestContext, plaintext: bytes) -> bytes: + return await self._current.encrypt(ctx, plaintext) + + async def decrypt(self, ctx: RequestContext, ciphertext: bytes) -> bytes: + last_exc: EncryptionError | None = None + for km in self._active_managers(): + try: + return await km.decrypt(ctx, ciphertext) + except EncryptionError as exc: + # Wrong key / sealed / tampered *for this manager* — try the next. + last_exc = exc + if last_exc is not None: + raise last_exc + raise KeyUnavailableError( # pragma: no cover - current is always tried + f"no key could decrypt for tenant {ctx.tenant_id!r}", + tenant_id=str(ctx.tenant_id), + ) + + async def rewrap(self, ctx: RequestContext, ciphertext: bytes) -> bytes: + """Re-encrypt a blob under the current key (background migration).""" + return await self._current.encrypt(ctx, await self.decrypt(ctx, ciphertext)) + + async def health(self) -> bool: + return await self._current.health() + + def _active_managers(self) -> list[KeyManager]: + # Current first (the common case), then retired keys that haven't expired. + now = self._clock() + managers = [self._current] + managers.extend( + r.key_manager for r in self._retired if r.expires_at is None or r.expires_at > now + ) + return managers diff --git a/packages/ragctl/src/ragctl/main.py b/packages/ragctl/src/ragctl/main.py index dd8c329..500e098 100644 --- a/packages/ragctl/src/ragctl/main.py +++ b/packages/ragctl/src/ragctl/main.py @@ -4889,17 +4889,20 @@ def audit( def kms( text: str = typer.Option("PII: alice@example.com", "--text", help="Sample plaintext."), tenant: str = typer.Option("acme", "--tenant", help="Tenant to encrypt for."), + rotate: bool = typer.Option(False, "--rotate", help="Also demo zero-downtime KEK rotation."), ) -> None: """Demo Step 6.7 BYOK envelope encryption with an in-process LocalKeyManager. Builds per-tenant KEKs, encrypts → decrypts (round-trip), then shows the two BYOK guarantees: per-tenant **isolation** (another tenant can't read the blob) - and **sealing** (a tenant with no KEK is denied with a typed error). No - external services / KMS. + and **sealing** (a tenant with no KEK is denied with a typed error). With + ``--rotate`` it also demos zero-downtime key rotation (6.7d): old + new data + both decrypt, ``rewrap`` migrates old blobs, and an expired retired key seals + old data. No external services / KMS. Example:: - ragctl kms --tenant acme --text "secret" + ragctl kms --tenant acme --text "secret" --rotate """ import asyncio import os @@ -4945,6 +4948,38 @@ async def _run() -> None: except KeyUnavailableError: typer.echo(" sealing: ok (tenant with no KEK is denied)") + if rotate: + from datetime import UTC, datetime, timedelta + + from rag_core.rotating_key_manager import RetiredKey, RotatingKeyManager + + v1 = LocalKeyManager(keks={tenant: os.urandom(32)}) + v2 = LocalKeyManager(keks={tenant: os.urandom(32)}) + old = await v1.encrypt(_ctx(tenant), pt) # wrapped before rotation + live = RotatingKeyManager( + current=v2, + retired=[RetiredKey(v1, expires_at=datetime.now(UTC) + timedelta(days=1))], + ) + new = await live.encrypt(_ctx(tenant), pt) + migrated = await live.rewrap(_ctx(tenant), old) + expired = RotatingKeyManager( + current=v2, + retired=[RetiredKey(v1, expires_at=datetime.now(UTC) - timedelta(seconds=1))], + ) + new_ok = await live.decrypt(_ctx(tenant), new) == pt + old_ok = await live.decrypt(_ctx(tenant), old) == pt + rewrap_ok = await v2.decrypt(_ctx(tenant), migrated) == pt + typer.echo("\nkey rotation — current=v2, retired=v1") + typer.echo("─" * 64) + typer.echo(f" new data (current): {'ok' if new_ok else 'FAIL'}") + typer.echo(f" old data (retired): {'ok' if old_ok else 'FAIL'}") + typer.echo(f" rewrap → current only: {'ok' if rewrap_ok else 'FAIL'}") + try: + await expired.decrypt(_ctx(tenant), old) + typer.echo(" expiry seals old data: FAIL") + except EncryptionError: + typer.echo(" expiry seals old data: ok (unreadable after retention)") + asyncio.run(_run()) diff --git a/packages/ragctl/tests/test_kms.py b/packages/ragctl/tests/test_kms.py index fa39927..460af72 100644 --- a/packages/ragctl/tests/test_kms.py +++ b/packages/ragctl/tests/test_kms.py @@ -18,3 +18,13 @@ def test_kms_demo_round_trip_isolation_sealing() -> None: assert "round-trip: ok" in result.output assert "isolation: ok" in result.output assert "sealing: ok" in result.output + + +def test_kms_rotate_demo() -> None: + result = runner.invoke(app, ["kms", "--tenant", "acme", "--text", "secret", "--rotate"]) + assert result.exit_code == 0, result.output + assert "key rotation" in result.output + assert "new data (current): ok" in result.output + assert "old data (retired): ok" in result.output + assert "rewrap → current only: ok" in result.output + assert "expiry seals old data: ok" in result.output diff --git a/tests/kms/test_rotation.py b/tests/kms/test_rotation.py new file mode 100644 index 0000000..e657c75 --- /dev/null +++ b/tests/kms/test_rotation.py @@ -0,0 +1,137 @@ +"""Zero-downtime KEK rotation — RotatingKeyManager (Step 6.7d). + +Exercises the rotation acceptance criteria over the real ``LocalKeyManager``: +after rotating, both old and new data decrypt; ``rewrap`` migrates old blobs to +the current key; once a retired key expires, old (un-rewrapped) data is sealed +with no impact on current-key data. +""" + +from __future__ import annotations + +import os +from datetime import UTC, datetime, timedelta + +import pytest +from rag_backends import LocalKeyManager +from rag_core.errors import EncryptionError, KeyUnavailableError +from rag_core.rotating_key_manager import RetiredKey, RotatingKeyManager +from rag_core.types import ( + Principal, + PrincipalId, + PrincipalKind, + RequestContext, + TenantId, +) + + +def _ctx(tenant: str = "acme") -> RequestContext: + tid = TenantId(tenant) + return RequestContext( + tenant_id=tid, + principal=Principal( + id=PrincipalId("p"), kind=PrincipalKind.user, display_name="p", tenant_id=tid + ), + ) + + +def _local() -> LocalKeyManager: + return LocalKeyManager(keks={"acme": os.urandom(32)}) + + +def _future() -> datetime: + return datetime.now(UTC) + timedelta(days=1) + + +def _past() -> datetime: + return datetime.now(UTC) - timedelta(seconds=1) + + +async def test_encrypt_uses_current_and_round_trips() -> None: + v1, v2 = _local(), _local() + rot = RotatingKeyManager(current=v2, retired=[RetiredKey(v1, expires_at=_future())]) + blob = await rot.encrypt(_ctx(), b"new data") + # encrypted with the current key → the current-only manager can read it + assert await v2.decrypt(_ctx(), blob) == b"new data" + assert await rot.decrypt(_ctx(), blob) == b"new data" + + +async def test_old_data_decrypts_via_retired_key() -> None: + v1, v2 = _local(), _local() + old = await v1.encrypt(_ctx(), b"old data") # wrapped before rotation + rot = RotatingKeyManager(current=v2, retired=[RetiredKey(v1, expires_at=_future())]) + assert await rot.decrypt(_ctx(), old) == b"old data" + + +async def test_rewrap_migrates_to_current() -> None: + v1, v2 = _local(), _local() + old = await v1.encrypt(_ctx(), b"old data") + rot = RotatingKeyManager(current=v2, retired=[RetiredKey(v1, expires_at=_future())]) + migrated = await rot.rewrap(_ctx(), old) + # after rewrap the blob is readable by the current key alone (retired no longer needed) + assert await v2.decrypt(_ctx(), migrated) == b"old data" + + +async def test_expired_retired_key_seals_old_data() -> None: + v1, v2 = _local(), _local() + old = await v1.encrypt(_ctx(), b"old data") + new = await v2.encrypt(_ctx(), b"new data") + rot = RotatingKeyManager(current=v2, retired=[RetiredKey(v1, expires_at=_past())]) + # current-key data is unaffected … + assert await rot.decrypt(_ctx(), new) == b"new data" + # … but old data wrapped under the now-expired key is unreadable (sealed) + with pytest.raises(EncryptionError): + await rot.decrypt(_ctx(), old) + + +async def test_tampered_blob_fails() -> None: + v1, v2 = _local(), _local() + rot = RotatingKeyManager(current=v2, retired=[RetiredKey(v1, expires_at=_future())]) + blob = bytearray(await rot.encrypt(_ctx(), b"new data")) + blob[-1] ^= 0x01 + with pytest.raises(EncryptionError): + await rot.decrypt(_ctx(), bytes(blob)) + + +async def test_no_key_at_all_is_key_unavailable() -> None: + # neither current nor retired has a KEK for this tenant → sealed at encrypt + rot = RotatingKeyManager(current=LocalKeyManager(keks={})) + with pytest.raises(KeyUnavailableError): + await rot.encrypt(_ctx("ghost"), b"x") + + +async def test_multiple_retired_generations_all_decrypt() -> None: + v1, v2, v3 = _local(), _local(), _local() + b1 = await v1.encrypt(_ctx(), b"gen1") + b2 = await v2.encrypt(_ctx(), b"gen2") + rot = RotatingKeyManager( + current=v3, + retired=[RetiredKey(v2, expires_at=_future()), RetiredKey(v1, expires_at=_future())], + ) + b3 = await rot.encrypt(_ctx(), b"gen3") + assert await rot.decrypt(_ctx(), b1) == b"gen1" + assert await rot.decrypt(_ctx(), b2) == b"gen2" + assert await rot.decrypt(_ctx(), b3) == b"gen3" + + +async def test_clock_injection_controls_expiry() -> None: + v1, v2 = _local(), _local() + old = await v1.encrypt(_ctx(), b"old data") + expiry = datetime(2026, 1, 1, tzinfo=UTC) + rot = RotatingKeyManager( + current=v2, + retired=[RetiredKey(v1, expires_at=expiry)], + clock=lambda: datetime(2025, 12, 31, tzinfo=UTC), # before expiry + ) + assert await rot.decrypt(_ctx(), old) == b"old data" + rot_after = RotatingKeyManager( + current=v2, + retired=[RetiredKey(v1, expires_at=expiry)], + clock=lambda: datetime(2026, 1, 2, tzinfo=UTC), # after expiry + ) + with pytest.raises(EncryptionError): + await rot_after.decrypt(_ctx(), old) + + +async def test_health_delegates_to_current() -> None: + rot = RotatingKeyManager(current=_local()) + assert await rot.health() is True