feat: rotate shared private-state keys (Closes #1237)

src/available_issue_snapshot_store.py

@@ -12,7 +12,7 @@ from src.private_state import connect_private_sqlite
 from src.state_encryption import (
     PrivateStateCipher,
     PrivateStateEncryptionError,
-    private_state_encryption_key,
+    private_state_encryption_config,
 )
 
 
@@ -34,7 +34,7 @@ class AvailableIssueSnapshotStore:
         self.path = Path(path)
         self.clock = clock or time.time
         self._cipher = PrivateStateCipher(
-            encryption_key if encryption_key is not None else private_state_encryption_key(),
+            encryption_key if encryption_key is not None else private_state_encryption_config(),
             store="available-issue-snapshot",
         )
         with self._connect() as connection:

src/idempotency.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from typing import Any, Callable
 
 from src.private_state import connect_private_sqlite
-from src.state_encryption import PrivateStateCipher, private_state_encryption_key
+from src.state_encryption import PrivateStateCipher, private_state_encryption_config
 
 
 @dataclass(frozen=True)
@@ -38,7 +38,7 @@ class IdempotencyLedger:
         self.lock_timeout_seconds = lock_timeout_seconds
         self.clock = clock
         self._cipher = PrivateStateCipher(
-            encryption_key if encryption_key is not None else private_state_encryption_key(),
+            encryption_key if encryption_key is not None else private_state_encryption_config(),
             store="idempotency-ledger",
         )
         self._initialize()

src/later_store.py

@@ -6,7 +6,7 @@ from datetime import datetime
 from pathlib import Path
 
 from src.private_state import connect_private_sqlite
-from src.state_encryption import PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_key
+from src.state_encryption import PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_config
 
 
 class LaterStore:
@@ -26,7 +26,7 @@ class LaterStore:
         self.operation_retention_seconds = operation_retention_seconds
         self.clock = clock
         self._cipher = PrivateStateCipher(
-            encryption_key if encryption_key is not None else private_state_encryption_key(),
+            encryption_key if encryption_key is not None else private_state_encryption_config(),
             store="later",
         )
         self._initialize()
@@ -124,7 +124,7 @@ class LaterStore:
             "SELECT item_id, revision FROM later_item_revisions WHERE login = ?",
             (login,),
         ):
-            if stored_item_id.startswith("v1:"):
+            if stored_item_id.startswith(("v1:", "v2:")):
                 item_id, _legacy = self._cipher.open(
                     stored_item_id, binding=f"item-revision:{login}"
                 )
@@ -141,7 +141,7 @@ class LaterStore:
     ) -> dict[str, str]:
         migrated = dict(stored_ids)
         for item_id, stored_item_id in stored_ids.items():
-            if stored_item_id.startswith("v1:"):
+            if stored_item_id.startswith(("v1:", "v2:")):
                 continue
             sealed_item_id = self._cipher.seal(
                 item_id, binding=f"item-revision:{login}"

src/live_snapshot_store.py

@@ -15,7 +15,7 @@ from src.private_state import connect_private_sqlite
 from src.state_encryption import (
     PrivateStateCipher,
     PrivateStateEncryptionError,
-    private_state_encryption_key,
+    private_state_encryption_config,
 )
 
 SECTIONS = ("context", "events", "notifications")
@@ -63,7 +63,7 @@ class LiveSnapshotStore:
         self.path = Path(path)
         self.clock = clock or time.time
         self._cipher = PrivateStateCipher(
-            encryption_key if encryption_key is not None else private_state_encryption_key(),
+            encryption_key if encryption_key is not None else private_state_encryption_config(),
             store="live-snapshot",
         )
         self._initialize()

src/saved_search_store.py

@@ -5,7 +5,7 @@ import sqlite3
 from pathlib import Path
 
 from src.private_state import connect_private_sqlite
-from src.state_encryption import PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_key
+from src.state_encryption import PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_config
 
 
 _REPOSITORY = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$")
@@ -33,7 +33,7 @@ class SavedSearchStore:
         self.limit = limit
         self.timeout = timeout
         self._cipher = PrivateStateCipher(
-            encryption_key if encryption_key is not None else private_state_encryption_key(),
+            encryption_key if encryption_key is not None else private_state_encryption_config(),
             store="saved-searches",
         )
         self._initialize()

src/security_event_store.py

@@ -10,7 +10,7 @@ from src.private_state import connect_private_sqlite
 from src.state_encryption import (
     PrivateStateCipher,
     PrivateStateEncryptionError,
-    private_state_encryption_key,
+    private_state_encryption_config,
 )
 
 
@@ -53,7 +53,7 @@ class SecurityEventStore:
         self.lock_timeout_seconds = lock_timeout_seconds
         try:
             self._cipher = PrivateStateCipher(
-                encryption_key if encryption_key is not None else private_state_encryption_key(),
+                encryption_key if encryption_key is not None else private_state_encryption_config(),
                 store="security-events",
             )
         except PrivateStateEncryptionError as exc:

src/state_encryption.py

@@ -6,6 +6,8 @@ import base64
 import binascii
 import json
 import os
+import re
+from collections.abc import Mapping
 
 from cryptography.exceptions import InvalidTag
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
@@ -16,7 +18,7 @@ class PrivateStateEncryptionError(RuntimeError):
 
 
 def decode_private_state_encryption_key(encoded: str) -> bytes:
-    """Decode the independently injected 256-bit private-state key."""
+    """Decode an independently injected 256-bit private-state key."""
     try:
         key = base64.b64decode(encoded, validate=True)
     except (binascii.Error, ValueError) as error:
@@ -31,38 +33,130 @@ def decode_private_state_encryption_key(encoded: str) -> bytes:
 
 
 def private_state_encryption_key() -> bytes:
+    """Load the legacy single key setting."""
     return decode_private_state_encryption_key(
         os.getenv("STACKCHAIN_PRIVATE_STATE_ENCRYPTION_KEY", "")
     )
 
 
+_KEY_ID = re.compile(r"^[A-Za-z0-9_-]{1,32}$")
+_MAX_KEYS = 4
+
+
+def decode_private_state_encryption_keyring(
+    encoded: str, active_key_id: str
+) -> tuple[dict[str, bytes], str]:
+    """Decode a bounded key-ID-to-key JSON object and validate its active key."""
+    try:
+        raw = json.loads(encoded)
+    except (TypeError, json.JSONDecodeError) as error:
+        raise PrivateStateEncryptionError(
+            "private state encryption keyring is invalid"
+        ) from error
+    if not isinstance(raw, dict) or not raw:
+        raise PrivateStateEncryptionError("private state encryption keyring is invalid")
+    if len(raw) > _MAX_KEYS:
+        raise PrivateStateEncryptionError(
+            f"private state encryption keyring supports at most {_MAX_KEYS} keys"
+        )
+    keys: dict[str, bytes] = {}
+    for key_id, value in raw.items():
+        if not isinstance(key_id, str) or not _KEY_ID.fullmatch(key_id):
+            raise PrivateStateEncryptionError("private state encryption key ID is invalid")
+        if not isinstance(value, str):
+            raise PrivateStateEncryptionError("private state encryption keyring is invalid")
+        keys[key_id] = decode_private_state_encryption_key(value)
+    if active_key_id not in keys:
+        raise PrivateStateEncryptionError(
+            "private state active key is missing from the keyring"
+        )
+    return keys, active_key_id
+
+
+def private_state_encryption_config() -> bytes | tuple[dict[str, bytes], str]:
+    """Load a rotation keyring when configured, otherwise the legacy single key."""
+    encoded = os.getenv("STACKCHAIN_PRIVATE_STATE_ENCRYPTION_KEYS", "")
+    active = os.getenv("STACKCHAIN_PRIVATE_STATE_ACTIVE_KEY_ID", "")
+    if encoded or active:
+        return decode_private_state_encryption_keyring(encoded, active)
+    return private_state_encryption_key()
+
+
 class PrivateStateCipher:
     """Seal JSON values with store-specific authenticated context."""
 
-    def __init__(self, key: bytes, *, store: str):
-        if not isinstance(key, bytes) or len(key) != 32:
-            raise PrivateStateEncryptionError(
-                "private state encryption requires exactly 32 key bytes"
-            )
+    def __init__(
+        self,
+        key: bytes | Mapping[str, bytes] | tuple[dict[str, bytes], str],
+        *,
+        store: str,
+        active_key_id: str | None = None,
+    ):
+        if isinstance(key, tuple):
+            key, active_key_id = key
+        if isinstance(key, Mapping):
+            if not key or len(key) > _MAX_KEYS or active_key_id not in key:
+                raise PrivateStateEncryptionError(
+                    "private state encryption keyring configuration is invalid"
+                )
+            ciphers: dict[str, AESGCM] = {}
+            for key_id, raw_key in key.items():
+                if not isinstance(key_id, str) or not _KEY_ID.fullmatch(key_id):
+                    raise PrivateStateEncryptionError(
+                        "private state encryption key ID is invalid"
+                    )
+                if not isinstance(raw_key, bytes) or len(raw_key) != 32:
+                    raise PrivateStateEncryptionError(
+                        "private state encryption requires exactly 32 key bytes"
+                    )
+                ciphers[key_id] = AESGCM(raw_key)
+            self._ciphers = ciphers
+            self._active_key_id = active_key_id
+            self._legacy_cipher = None
+        else:
+            if not isinstance(key, bytes) or len(key) != 32:
+                raise PrivateStateEncryptionError(
+                    "private state encryption requires exactly 32 key bytes"
+                )
+            self._ciphers = {}
+            self._active_key_id = None
+            self._legacy_cipher = AESGCM(key)
         if not store or "\0" in store:
             raise ValueError("private state store identity is invalid")
-        self._cipher = AESGCM(key)
         self._store = store
 
-    def _aad(self, binding: str) -> bytes:
-        return f"stackchain:private-state:v1\0{self._store}\0{binding}".encode()
+    def _aad(self, binding: str, *, version: str = "v1", key_id: str = "") -> bytes:
+        prefix = f"stackchain:private-state:{version}\0{self._store}\0{binding}"
+        if version == "v2":
+            prefix += f"\0{key_id}"
+        return prefix.encode()
 
     def seal(self, value: object, *, binding: str = "singleton") -> str:
         plaintext = json.dumps(value, separators=(",", ":")).encode()
         nonce = os.urandom(12)
-        sealed = nonce + self._cipher.encrypt(nonce, plaintext, self._aad(binding))
+        if self._active_key_id is not None:
+            key_id = self._active_key_id
+            sealed = nonce + self._ciphers[key_id].encrypt(
+                nonce, plaintext, self._aad(binding, version="v2", key_id=key_id)
+            )
+            return f"v2:{key_id}:" + base64.urlsafe_b64encode(sealed).decode()
+        sealed = nonce + self._legacy_cipher.encrypt(
+            nonce, plaintext, self._aad(binding)
+        )
         return "v1:" + base64.urlsafe_b64encode(sealed).decode()
 
+    @staticmethod
+    def _decode_sealed(payload: str) -> bytes:
+        sealed = base64.b64decode(payload, altchars=b"-_", validate=True)
+        if len(sealed) < 28:
+            raise ValueError("encrypted payload is too short")
+        return sealed
+
     def open(self, payload: str, *, binding: str = "singleton") -> tuple[object, bool]:
-        """Return the decoded value and whether plaintext migration is required."""
+        """Return the decoded value and whether active-key rewrapping is required."""
         if not isinstance(payload, str):
             raise PrivateStateEncryptionError("private state could not be decrypted")
-        if not payload.startswith("v1:"):
+        if not payload.startswith(("v1:", "v2:")):
             try:
                 return json.loads(payload), True
             except (TypeError, json.JSONDecodeError) as error:
@@ -70,13 +164,37 @@ class PrivateStateCipher:
                     "private state could not be decrypted"
                 ) from error
         try:
-            sealed = base64.b64decode(payload[3:], altchars=b"-_", validate=True)
-            if len(sealed) < 28:
-                raise ValueError("encrypted payload is too short")
-            plaintext = self._cipher.decrypt(
-                sealed[:12], sealed[12:], self._aad(binding)
+            if payload.startswith("v2:"):
+                _, key_id, encoded = payload.split(":", 2)
+                cipher = self._ciphers.get(key_id)
+                if cipher is None:
+                    raise ValueError("encrypted payload uses an unavailable key")
+                sealed = self._decode_sealed(encoded)
+                plaintext = cipher.decrypt(
+                    sealed[:12],
+                    sealed[12:],
+                    self._aad(binding, version="v2", key_id=key_id),
+                )
+                return json.loads(plaintext), key_id != self._active_key_id
+
+            sealed = self._decode_sealed(payload[3:])
+            candidates = (
+                [self._legacy_cipher]
+                if self._legacy_cipher is not None
+                else list(self._ciphers.values())
             )
-            return json.loads(plaintext), False
+            plaintext = None
+            for cipher in candidates:
+                try:
+                    plaintext = cipher.decrypt(
+                        sealed[:12], sealed[12:], self._aad(binding)
+                    )
+                    break
+                except InvalidTag:
+                    continue
+            if plaintext is None:
+                raise InvalidTag
+            return json.loads(plaintext), self._active_key_id is not None
         except (
             binascii.Error,
             InvalidTag,

src/today_store.py

@@ -8,7 +8,7 @@ from pathlib import Path
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
 from src.private_state import connect_private_sqlite
-from src.state_encryption import PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_key
+from src.state_encryption import PrivateStateCipher, PrivateStateEncryptionError, private_state_encryption_config
 
 
 class TodayPlanFull(ValueError):
@@ -77,7 +77,7 @@ class TodayStore:
         self.recap_limit = recap_limit
         self.clock = clock
         self._cipher = PrivateStateCipher(
-            encryption_key if encryption_key is not None else private_state_encryption_key(),
+            encryption_key if encryption_key is not None else private_state_encryption_config(),
             store="today",
         )
         self._initialize()
@@ -228,7 +228,7 @@ class TodayStore:
     def _snapshot(self, row, login: str) -> tuple[dict, bool]:
         if row is None:
             return {"revision": 0, "ids": [], "capacity_minutes": None, "estimates": {}}, False
-        if row[1].startswith("v1:"):
+        if row[1].startswith(("v1:", "v2:")):
             payload, legacy = self._cipher.open(row[1], binding=f"plan:{login}")
             if not isinstance(payload, dict):
                 raise PrivateStateEncryptionError("private state could not be decrypted")
@@ -828,7 +828,7 @@ class TodayStore:
     def _session_snapshot(self, row, login: str) -> tuple[dict, bool]:
         if row is None:
             return self._empty_session(), False
-        if row[1].startswith("v1:"):
+        if row[1].startswith(("v1:", "v2:")):
             payload, legacy = self._cipher.open(row[1], binding=f"session:{login}")
             if not isinstance(payload, dict):
                 raise PrivateStateEncryptionError("private state could not be decrypted")
@@ -929,7 +929,7 @@ class TodayStore:
     def _recap_snapshot(
         self, login: str, encrypted_session_id: str, created_at: float, serialized: str
     ) -> tuple[dict, bool]:
-        if encrypted_session_id.startswith("v1:"):
+        if encrypted_session_id.startswith(("v1:", "v2:")):
             session_id, legacy_session = self._cipher.open(
                 encrypted_session_id, binding=f"recap-id:{login}"
             )
@@ -1021,7 +1021,7 @@ class TodayStore:
 
     def _time_log_snapshot(self, login: str, row) -> tuple[dict, bool]:
         encrypted_session, encrypted_identity, payload_value, status_value = row
-        if encrypted_session.startswith("v1:"):
+        if encrypted_session.startswith(("v1:", "v2:")):
             session_id = self._cipher.open(
                 encrypted_session, binding=f"time-log-session:{login}"
             )[0]