Add new SSE, KMS encryptions
This commit is contained in:
343
app/kms.py
Normal file
343
app/kms.py
Normal file
@@ -0,0 +1,343 @@
|
||||
"""Key Management Service (KMS) for encryption key management."""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import secrets
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
|
||||
from .encryption import EncryptionError, EncryptionProvider, EncryptionResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class KMSKey:
|
||||
"""Represents a KMS encryption key."""
|
||||
key_id: str
|
||||
description: str
|
||||
created_at: str
|
||||
enabled: bool = True
|
||||
key_material: bytes = field(default_factory=lambda: b"", repr=False)
|
||||
|
||||
@property
|
||||
def arn(self) -> str:
|
||||
return f"arn:aws:kms:local:000000000000:key/{self.key_id}"
|
||||
|
||||
def to_dict(self, include_key: bool = False) -> Dict[str, Any]:
|
||||
data = {
|
||||
"KeyId": self.key_id,
|
||||
"Arn": self.arn,
|
||||
"Description": self.description,
|
||||
"CreationDate": self.created_at,
|
||||
"Enabled": self.enabled,
|
||||
"KeyState": "Enabled" if self.enabled else "Disabled",
|
||||
"KeyUsage": "ENCRYPT_DECRYPT",
|
||||
"KeySpec": "SYMMETRIC_DEFAULT",
|
||||
}
|
||||
if include_key:
|
||||
data["KeyMaterial"] = base64.b64encode(self.key_material).decode()
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "KMSKey":
|
||||
key_material = b""
|
||||
if "KeyMaterial" in data:
|
||||
key_material = base64.b64decode(data["KeyMaterial"])
|
||||
return cls(
|
||||
key_id=data["KeyId"],
|
||||
description=data.get("Description", ""),
|
||||
created_at=data.get("CreationDate", datetime.now(timezone.utc).isoformat()),
|
||||
enabled=data.get("Enabled", True),
|
||||
key_material=key_material,
|
||||
)
|
||||
|
||||
|
||||
class KMSEncryptionProvider(EncryptionProvider):
|
||||
"""Encryption provider using a specific KMS key."""
|
||||
|
||||
def __init__(self, kms: "KMSManager", key_id: str):
|
||||
self.kms = kms
|
||||
self.key_id = key_id
|
||||
|
||||
@property
|
||||
def KEY_ID(self) -> str:
|
||||
return self.key_id
|
||||
|
||||
def generate_data_key(self) -> tuple[bytes, bytes]:
|
||||
"""Generate a data key encrypted with the KMS key."""
|
||||
return self.kms.generate_data_key(self.key_id)
|
||||
|
||||
def encrypt(self, plaintext: bytes, context: Dict[str, str] | None = None) -> EncryptionResult:
|
||||
"""Encrypt data using envelope encryption with KMS."""
|
||||
data_key, encrypted_data_key = self.generate_data_key()
|
||||
|
||||
aesgcm = AESGCM(data_key)
|
||||
nonce = secrets.token_bytes(12)
|
||||
ciphertext = aesgcm.encrypt(nonce, plaintext,
|
||||
json.dumps(context).encode() if context else None)
|
||||
|
||||
return EncryptionResult(
|
||||
ciphertext=ciphertext,
|
||||
nonce=nonce,
|
||||
key_id=self.key_id,
|
||||
encrypted_data_key=encrypted_data_key,
|
||||
)
|
||||
|
||||
def decrypt(self, ciphertext: bytes, nonce: bytes, encrypted_data_key: bytes,
|
||||
key_id: str, context: Dict[str, str] | None = None) -> bytes:
|
||||
"""Decrypt data using envelope encryption with KMS."""
|
||||
data_key = self.kms.decrypt_data_key(key_id, encrypted_data_key, context)
|
||||
|
||||
aesgcm = AESGCM(data_key)
|
||||
try:
|
||||
return aesgcm.decrypt(nonce, ciphertext,
|
||||
json.dumps(context).encode() if context else None)
|
||||
except Exception as exc:
|
||||
raise EncryptionError(f"Failed to decrypt data: {exc}") from exc
|
||||
|
||||
|
||||
class KMSManager:
|
||||
"""Manages KMS keys and operations.
|
||||
|
||||
This is a local implementation that mimics AWS KMS functionality.
|
||||
Keys are stored encrypted on disk.
|
||||
"""
|
||||
|
||||
def __init__(self, keys_path: Path, master_key_path: Path):
|
||||
self.keys_path = keys_path
|
||||
self.master_key_path = master_key_path
|
||||
self._keys: Dict[str, KMSKey] = {}
|
||||
self._master_key: bytes | None = None
|
||||
self._loaded = False
|
||||
|
||||
@property
|
||||
def master_key(self) -> bytes:
|
||||
"""Load or create the master key for encrypting KMS keys."""
|
||||
if self._master_key is None:
|
||||
if self.master_key_path.exists():
|
||||
self._master_key = base64.b64decode(
|
||||
self.master_key_path.read_text().strip()
|
||||
)
|
||||
else:
|
||||
self._master_key = secrets.token_bytes(32)
|
||||
self.master_key_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.master_key_path.write_text(
|
||||
base64.b64encode(self._master_key).decode()
|
||||
)
|
||||
return self._master_key
|
||||
|
||||
def _load_keys(self) -> None:
|
||||
"""Load keys from disk."""
|
||||
if self._loaded:
|
||||
return
|
||||
|
||||
if self.keys_path.exists():
|
||||
try:
|
||||
data = json.loads(self.keys_path.read_text(encoding="utf-8"))
|
||||
for key_data in data.get("keys", []):
|
||||
key = KMSKey.from_dict(key_data)
|
||||
if key_data.get("EncryptedKeyMaterial"):
|
||||
encrypted = base64.b64decode(key_data["EncryptedKeyMaterial"])
|
||||
key.key_material = self._decrypt_key_material(encrypted)
|
||||
self._keys[key.key_id] = key
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self._loaded = True
|
||||
|
||||
def _save_keys(self) -> None:
|
||||
"""Save keys to disk (with encrypted key material)."""
|
||||
keys_data = []
|
||||
for key in self._keys.values():
|
||||
data = key.to_dict(include_key=False)
|
||||
encrypted = self._encrypt_key_material(key.key_material)
|
||||
data["EncryptedKeyMaterial"] = base64.b64encode(encrypted).decode()
|
||||
keys_data.append(data)
|
||||
|
||||
self.keys_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.keys_path.write_text(
|
||||
json.dumps({"keys": keys_data}, indent=2),
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
def _encrypt_key_material(self, key_material: bytes) -> bytes:
|
||||
"""Encrypt key material with the master key."""
|
||||
aesgcm = AESGCM(self.master_key)
|
||||
nonce = secrets.token_bytes(12)
|
||||
ciphertext = aesgcm.encrypt(nonce, key_material, None)
|
||||
return nonce + ciphertext
|
||||
|
||||
def _decrypt_key_material(self, encrypted: bytes) -> bytes:
|
||||
"""Decrypt key material with the master key."""
|
||||
aesgcm = AESGCM(self.master_key)
|
||||
nonce = encrypted[:12]
|
||||
ciphertext = encrypted[12:]
|
||||
return aesgcm.decrypt(nonce, ciphertext, None)
|
||||
|
||||
def create_key(self, description: str = "", key_id: str | None = None) -> KMSKey:
|
||||
"""Create a new KMS key."""
|
||||
self._load_keys()
|
||||
|
||||
if key_id is None:
|
||||
key_id = str(uuid.uuid4())
|
||||
|
||||
if key_id in self._keys:
|
||||
raise EncryptionError(f"Key already exists: {key_id}")
|
||||
|
||||
key = KMSKey(
|
||||
key_id=key_id,
|
||||
description=description,
|
||||
created_at=datetime.now(timezone.utc).isoformat(),
|
||||
enabled=True,
|
||||
key_material=secrets.token_bytes(32),
|
||||
)
|
||||
|
||||
self._keys[key_id] = key
|
||||
self._save_keys()
|
||||
return key
|
||||
|
||||
def get_key(self, key_id: str) -> KMSKey | None:
|
||||
"""Get a key by ID."""
|
||||
self._load_keys()
|
||||
return self._keys.get(key_id)
|
||||
|
||||
def list_keys(self) -> List[KMSKey]:
|
||||
"""List all keys."""
|
||||
self._load_keys()
|
||||
return list(self._keys.values())
|
||||
|
||||
def enable_key(self, key_id: str) -> None:
|
||||
"""Enable a key."""
|
||||
self._load_keys()
|
||||
key = self._keys.get(key_id)
|
||||
if not key:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
key.enabled = True
|
||||
self._save_keys()
|
||||
|
||||
def disable_key(self, key_id: str) -> None:
|
||||
"""Disable a key."""
|
||||
self._load_keys()
|
||||
key = self._keys.get(key_id)
|
||||
if not key:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
key.enabled = False
|
||||
self._save_keys()
|
||||
|
||||
def delete_key(self, key_id: str) -> None:
|
||||
"""Delete a key (schedule for deletion in real KMS)."""
|
||||
self._load_keys()
|
||||
if key_id not in self._keys:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
del self._keys[key_id]
|
||||
self._save_keys()
|
||||
|
||||
def encrypt(self, key_id: str, plaintext: bytes,
|
||||
context: Dict[str, str] | None = None) -> bytes:
|
||||
"""Encrypt data directly with a KMS key."""
|
||||
self._load_keys()
|
||||
key = self._keys.get(key_id)
|
||||
if not key:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
if not key.enabled:
|
||||
raise EncryptionError(f"Key is disabled: {key_id}")
|
||||
|
||||
aesgcm = AESGCM(key.key_material)
|
||||
nonce = secrets.token_bytes(12)
|
||||
aad = json.dumps(context).encode() if context else None
|
||||
ciphertext = aesgcm.encrypt(nonce, plaintext, aad)
|
||||
|
||||
key_id_bytes = key_id.encode("utf-8")
|
||||
return len(key_id_bytes).to_bytes(2, "big") + key_id_bytes + nonce + ciphertext
|
||||
|
||||
def decrypt(self, ciphertext: bytes,
|
||||
context: Dict[str, str] | None = None) -> tuple[bytes, str]:
|
||||
"""Decrypt data directly with a KMS key.
|
||||
|
||||
Returns:
|
||||
Tuple of (plaintext, key_id)
|
||||
"""
|
||||
self._load_keys()
|
||||
|
||||
key_id_len = int.from_bytes(ciphertext[:2], "big")
|
||||
key_id = ciphertext[2:2 + key_id_len].decode("utf-8")
|
||||
rest = ciphertext[2 + key_id_len:]
|
||||
|
||||
key = self._keys.get(key_id)
|
||||
if not key:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
if not key.enabled:
|
||||
raise EncryptionError(f"Key is disabled: {key_id}")
|
||||
|
||||
nonce = rest[:12]
|
||||
encrypted = rest[12:]
|
||||
|
||||
aesgcm = AESGCM(key.key_material)
|
||||
aad = json.dumps(context).encode() if context else None
|
||||
try:
|
||||
plaintext = aesgcm.decrypt(nonce, encrypted, aad)
|
||||
return plaintext, key_id
|
||||
except Exception as exc:
|
||||
raise EncryptionError(f"Decryption failed: {exc}") from exc
|
||||
|
||||
def generate_data_key(self, key_id: str,
|
||||
context: Dict[str, str] | None = None) -> tuple[bytes, bytes]:
|
||||
"""Generate a data key and return both plaintext and encrypted versions.
|
||||
|
||||
Returns:
|
||||
Tuple of (plaintext_key, encrypted_key)
|
||||
"""
|
||||
self._load_keys()
|
||||
key = self._keys.get(key_id)
|
||||
if not key:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
if not key.enabled:
|
||||
raise EncryptionError(f"Key is disabled: {key_id}")
|
||||
|
||||
plaintext_key = secrets.token_bytes(32)
|
||||
|
||||
encrypted_key = self.encrypt(key_id, plaintext_key, context)
|
||||
|
||||
return plaintext_key, encrypted_key
|
||||
|
||||
def decrypt_data_key(self, key_id: str, encrypted_key: bytes,
|
||||
context: Dict[str, str] | None = None) -> bytes:
|
||||
"""Decrypt a data key."""
|
||||
plaintext, _ = self.decrypt(encrypted_key, context)
|
||||
return plaintext
|
||||
|
||||
def get_provider(self, key_id: str | None = None) -> KMSEncryptionProvider:
|
||||
"""Get an encryption provider for a specific key."""
|
||||
self._load_keys()
|
||||
|
||||
if key_id is None:
|
||||
if not self._keys:
|
||||
key = self.create_key("Default KMS Key")
|
||||
key_id = key.key_id
|
||||
else:
|
||||
key_id = next(iter(self._keys.keys()))
|
||||
|
||||
if key_id not in self._keys:
|
||||
raise EncryptionError(f"Key not found: {key_id}")
|
||||
|
||||
return KMSEncryptionProvider(self, key_id)
|
||||
|
||||
def re_encrypt(self, ciphertext: bytes, destination_key_id: str,
|
||||
source_context: Dict[str, str] | None = None,
|
||||
destination_context: Dict[str, str] | None = None) -> bytes:
|
||||
"""Re-encrypt data with a different key."""
|
||||
|
||||
plaintext, source_key_id = self.decrypt(ciphertext, source_context)
|
||||
|
||||
return self.encrypt(destination_key_id, plaintext, destination_context)
|
||||
|
||||
def generate_random(self, num_bytes: int = 32) -> bytes:
|
||||
"""Generate cryptographically secure random bytes."""
|
||||
if num_bytes < 1 or num_bytes > 1024:
|
||||
raise EncryptionError("Number of bytes must be between 1 and 1024")
|
||||
return secrets.token_bytes(num_bytes)
|
||||
Reference in New Issue
Block a user