stackchain-dashboard/src/state_encryption.py
timmy 965a73ab49
All checks were successful
CI / lint (pull_request) Successful in 3m24s
CI / build-release (pull_request) Successful in 6s
CI / browser-journey (pull_request) Successful in 5m2s
CI / release-candidate (pull_request) Has been skipped
feat: rotate shared private-state keys (Closes #1237)
2026-08-21 21:12:54 +00:00

208 lines
8.1 KiB
Python

"""Authenticated envelopes for retained private dashboard state."""
from __future__ import annotations
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
class PrivateStateEncryptionError(RuntimeError):
"""Private retained state could not be authenticated or decrypted."""
def decode_private_state_encryption_key(encoded: str) -> bytes:
"""Decode an independently injected 256-bit private-state key."""
try:
key = base64.b64decode(encoded, validate=True)
except (binascii.Error, ValueError) as error:
raise PrivateStateEncryptionError(
"private state encryption key is invalid"
) from error
if len(key) != 32:
raise PrivateStateEncryptionError(
"private state encryption key must decode to exactly 32 bytes"
)
return key
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 | 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._store = store
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)
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 active-key rewrapping is required."""
if not isinstance(payload, str):
raise PrivateStateEncryptionError("private state could not be decrypted")
if not payload.startswith(("v1:", "v2:")):
try:
return json.loads(payload), True
except (TypeError, json.JSONDecodeError) as error:
raise PrivateStateEncryptionError(
"private state could not be decrypted"
) from error
try:
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())
)
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,
UnicodeDecodeError,
ValueError,
json.JSONDecodeError,
) as error:
raise PrivateStateEncryptionError(
"private state could not be decrypted"
) from error