"""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