kqjy/MyFSIO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 33 Wiki Activity

Optimize KMS: cache AESGCM instance, remove duplicate get_provider

Browse Source
This commit is contained in:
kqjy
2026-02-09 17:01:19 +08:00
parent 54705ab9c4
commit 6e6d6d32bf
1 changed files with 7 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
app/kms.py
View File

@@ -160,6 +160,7 @@ class KMSManager:
self.generate_data_key_max_bytes = generate_data_key_max_bytes
self._keys: Dict[str, KMSKey] = {}
self._master_key: bytes | None = None
self._master_aesgcm: AESGCM | None = None
self._loaded = False
@property
@@ -191,6 +192,7 @@ class KMSManager:
msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
else:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
self._master_aesgcm = AESGCM(self._master_key)
return self._master_key
def _load_keys(self) -> None:
@@ -231,18 +233,16 @@ class KMSManager:
_set_secure_file_permissions(self.keys_path)
def _encrypt_key_material(self, key_material: bytes) -> bytes:
"""Encrypt key material with the master key."""
aesgcm = AESGCM(self.master_key)
_ = self.master_key
nonce = secrets.token_bytes(12)
ciphertext = aesgcm.encrypt(nonce, key_material, None)
ciphertext = self._master_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)
_ = self.master_key
nonce = encrypted[:12]
ciphertext = encrypted[12:]
return aesgcm.decrypt(nonce, ciphertext, None)
return self._master_aesgcm.decrypt(nonce, ciphertext, None)
def create_key(self, description: str = "", key_id: str | None = None) -> KMSKey:
"""Create a new KMS key."""
@@ -404,22 +404,6 @@ class KMSManager:
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: