feat(03-01): DB migrations, models, encryption service, and test scaffolds
- Add stripe and cryptography to shared pyproject.toml - Add recharts, @stripe/stripe-js, stripe to portal package.json (submodule) - Add billing fields to Tenant model (stripe_customer_id, subscription_status, agent_quota, trial_ends_at) - Add budget_limit_usd to Agent model - Create TenantLlmKey and StripeEvent models in billing.py (AuditBase and Base respectively) - Create KeyEncryptionService (MultiFernet encrypt/decrypt/rotate) in crypto.py - Create compute_budget_status helper in usage.py (threshold logic: ok/warning/exceeded) - Add platform_encryption_key, stripe_, slack_oauth settings to config.py - Create Alembic migration 005 with all schema changes, RLS, grants, and composite index - All 12 tests passing (key encryption roundtrip, rotation, budget thresholds)
This commit is contained in:
87
packages/shared/shared/crypto.py
Normal file
87
packages/shared/shared/crypto.py
Normal file
@@ -0,0 +1,87 @@
|
||||
"""
|
||||
KeyEncryptionService — Fernet-based encryption for BYO API keys.
|
||||
|
||||
Uses MultiFernet to support key rotation:
|
||||
- primary key: active encryption key (all new values encrypted with this)
|
||||
- previous key: optional previous key (supports decryption during rotation window)
|
||||
|
||||
The PLATFORM_ENCRYPTION_KEY environment variable must be a valid URL-safe
|
||||
base64-encoded 32-byte key, as generated by `Fernet.generate_key()`.
|
||||
|
||||
Usage:
|
||||
from shared.crypto import KeyEncryptionService
|
||||
svc = KeyEncryptionService(primary_key=settings.platform_encryption_key)
|
||||
ciphertext = svc.encrypt("sk-my-secret-key")
|
||||
plaintext = svc.decrypt(ciphertext)
|
||||
new_cipher = svc.rotate(old_ciphertext) # re-encrypts with primary key
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from cryptography.fernet import Fernet, MultiFernet
|
||||
|
||||
|
||||
class KeyEncryptionService:
|
||||
"""
|
||||
Encrypt and decrypt BYO API keys using Fernet symmetric encryption.
|
||||
|
||||
Fernet guarantees:
|
||||
- AES-128-CBC with PKCS7 padding
|
||||
- HMAC-SHA256 authentication
|
||||
- Random IV per encryption call (produces different ciphertext each time)
|
||||
- Timestamp in token (can enforce TTL if desired)
|
||||
|
||||
MultiFernet supports key rotation:
|
||||
- Encryption always uses the first (primary) key
|
||||
- Decryption tries all keys in order until one succeeds
|
||||
- rotate() decrypts with any key, re-encrypts with the primary key
|
||||
"""
|
||||
|
||||
def __init__(self, primary_key: str, previous_key: str = "") -> None:
|
||||
"""
|
||||
Initialise the service with one or two Fernet keys.
|
||||
|
||||
Args:
|
||||
primary_key: Active key for encryption and decryption. Must be a
|
||||
URL-safe base64-encoded 32-byte value (Fernet key).
|
||||
previous_key: Optional previous key retained only for decryption
|
||||
during a rotation window. Pass "" to omit.
|
||||
"""
|
||||
keys: list[Fernet] = [Fernet(primary_key.encode())]
|
||||
if previous_key:
|
||||
keys.append(Fernet(previous_key.encode()))
|
||||
self._multi = MultiFernet(keys)
|
||||
|
||||
def encrypt(self, plaintext: str) -> str:
|
||||
"""
|
||||
Encrypt a plaintext string.
|
||||
|
||||
Returns a URL-safe base64-encoded Fernet token (str).
|
||||
Calling encrypt() twice with the same plaintext produces different
|
||||
ciphertexts due to the random IV embedded in each Fernet token.
|
||||
"""
|
||||
return self._multi.encrypt(plaintext.encode()).decode()
|
||||
|
||||
def decrypt(self, ciphertext: str) -> str:
|
||||
"""
|
||||
Decrypt a Fernet token back to the original plaintext.
|
||||
|
||||
Raises:
|
||||
cryptography.fernet.InvalidToken: if the ciphertext is invalid,
|
||||
tampered, or cannot be decrypted by any of the known keys.
|
||||
"""
|
||||
return self._multi.decrypt(ciphertext.encode()).decode()
|
||||
|
||||
def rotate(self, ciphertext: str) -> str:
|
||||
"""
|
||||
Re-encrypt an existing ciphertext with the current primary key.
|
||||
|
||||
Useful for key rotation: after adding a new primary key and keeping
|
||||
the old key as previous_key, call rotate() on each stored ciphertext
|
||||
to migrate it to the new key. Once all values are rotated, the old
|
||||
key can be removed.
|
||||
|
||||
Returns a new Fernet token encrypted with the primary key.
|
||||
Raises InvalidToken if the ciphertext cannot be decrypted.
|
||||
"""
|
||||
return self._multi.rotate(ciphertext.encode()).decode()
|
||||
Reference in New Issue
Block a user