"""
Wallet Master Key Utilities
Shared logic for lazily loading, generating, and persisting
AES-256-GCM master keys in Redis for wallet encryption.
"""
import os
import base64
import logging
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
logger = logging.getLogger(__name__)
[docs]
async def ensure_master_key(
current_key: bytes | None,
redis_client,
redis_key: str,
env_var: str,
) -> bytes:
"""Return a 32-byte AES master key, loading or generating as needed.
Resolution order:
1. *current_key* – already loaded in memory (fast path).
2. Environment variable (*env_var*) – decode, validate, persist to Redis.
3. Redis (*redis_key*) – read previously persisted key.
4. Generate a new random key, persist to Redis.
Args:
current_key: The key already held in memory, or ``None``.
redis_client: An async Redis client.
redis_key: Redis key under which the master key is stored.
env_var: Name of the environment variable to check.
Returns:
A 32-byte ``bytes`` master key.
"""
# 1. Already loaded
if current_key is not None:
return current_key
# 2. Environment variable
key_b64 = os.environ.get(env_var)
if key_b64:
try:
key = base64.urlsafe_b64decode(key_b64)
if len(key) == 32:
# Persist to Redis so future restarts without the env var still work
await redis_client.set(redis_key, base64.urlsafe_b64encode(key).decode())
logger.info("Loaded %s from environment and persisted to Redis", env_var)
return key
else:
logger.warning("%s is not 32 bytes, ignoring", env_var)
except Exception as e:
logger.warning("Failed to decode %s: %s", env_var, e)
# 3. Redis
stored = await redis_client.get(redis_key)
if stored:
try:
if isinstance(stored, bytes):
stored = stored.decode()
key = base64.urlsafe_b64decode(stored)
if len(key) == 32:
logger.info("Loaded wallet master key from Redis (%s)", redis_key)
return key
else:
logger.warning("Redis key %s is not 32 bytes, regenerating", redis_key)
except Exception as e:
logger.warning("Failed to decode key from Redis (%s): %s", redis_key, e)
# 4. Generate and persist
key = AESGCM.generate_key(bit_length=256)
await redis_client.set(redis_key, base64.urlsafe_b64encode(key).decode())
logger.info(
"Generated new wallet master key and persisted to Redis (%s)", redis_key,
)
return key