Skip to content

Commit fedfb84

Browse files
aidehnaidehn
authored
KMS: On Demand Key Rotation for Imported Key Material (#13363)
1 parent dba528e commit fedfb84

File tree

12 files changed

+1107
-80
lines changed

12 files changed

+1107
-80
lines changed

‎localstack-core/localstack/services/kms/models.py‎

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,7 @@ class KmsCryptoKey:
173173
public_key: bytes | None
174174
private_key: bytes | None
175175
key_material: bytes
176+
pending_key_material: bytes | None
176177
key_spec: str
177178

178179
@staticmethod
@@ -217,6 +218,7 @@ def raise_validation():
217218
def __init__(self, key_spec: str, key_material: bytes | None = None):
218219
self.private_key = None
219220
self.public_key = None
221+
self.pending_key_material = None
220222
# Technically, key_material, being a symmetric encryption key, is only relevant for
221223
# key_spec == SYMMETRIC_DEFAULT.
222224
# But LocalStack uses symmetric encryption with this key_material even for other specs. Asymmetric keys are
@@ -248,8 +250,9 @@ def __init__(self, key_spec: str, key_material: bytes | None = None):
248250
self._serialize_key(key)
249251

250252
def load_key_material(self, material: bytes):
251-
if self.key_spec in [
252-
KeySpec.SYMMETRIC_DEFAULT,
253+
if self.key_spec == KeySpec.SYMMETRIC_DEFAULT:
254+
self.pending_key_material = material
255+
elif self.key_spec in [
253256
KeySpec.HMAC_224,
254257
KeySpec.HMAC_256,
255258
KeySpec.HMAC_384,
@@ -323,9 +326,28 @@ def __init__(
323326
# remove the _custom_key_material_ tag from the tags to not readily expose the custom key material
324327
del self.tags[TAG_KEY_CUSTOM_KEY_MATERIAL]
325328
self.crypto_key = KmsCryptoKey(self.metadata.get("KeySpec"), custom_key_material)
329+
self._internal_key_id = uuid.uuid4()
330+
331+
# The KMS implementation always provides a crypto key with key material which doesn't suit scenarios where a
332+
# KMS Key may have no key material e.g. for external keys. Don't expose the CurrentKeyMaterialId in those cases.
333+
if custom_key_material or (
334+
self.metadata["Origin"] == "AWS_KMS"
335+
and self.metadata["KeySpec"] == KeySpec.SYMMETRIC_DEFAULT
336+
):
337+
self.metadata["CurrentKeyMaterialId"] = self.generate_key_material_id(
338+
self.crypto_key.key_material
339+
)
340+
326341
self.rotation_period_in_days = 365
327342
self.next_rotation_date = None
328343

344+
def generate_key_material_id(self, key_material: bytes) -> str:
345+
# The KeyMaterialId depends on the key material and the KeyId. Use an internal ID to prevent brute forcing
346+
# the value of the key material from the public KeyId and KeyMaterialId.
347+
# https://docs.aws.amazon.com/kms/latest/APIReference/API_ImportKeyMaterial.html
348+
key_material_id_hex = uuid.uuid5(self._internal_key_id, key_material).hex
349+
return str(key_material_id_hex) * 2
350+
329351
def calculate_and_set_arn(self, account_id, region):
330352
self.metadata["Arn"] = kms_key_arn(self.metadata.get("KeyId"), account_id, region)
331353

@@ -746,8 +768,16 @@ def rotate_key_on_demand(self):
746768
f"The on-demand rotations limit has been reached for the given keyId. "
747769
f"No more on-demand rotations can be performed for this key: {self.metadata['Arn']}"
748770
)
749-
self.previous_keys.append(self.crypto_key.key_material)
750-
self.crypto_key = KmsCryptoKey(KeySpec.SYMMETRIC_DEFAULT)
771+
current_key_material = self.crypto_key.key_material
772+
pending_key_material = self.crypto_key.pending_key_material
773+
774+
self.previous_keys.append(current_key_material)
775+
776+
# If there is no pending material stored on the key, then key material will be generated.
777+
self.crypto_key = KmsCryptoKey(KeySpec.SYMMETRIC_DEFAULT, pending_key_material)
778+
self.metadata["CurrentKeyMaterialId"] = self.generate_key_material_id(
779+
self.crypto_key.key_material
780+
)
751781

752782

753783
class KmsGrant:

‎localstack-core/localstack/services/kms/provider.py‎

Lines changed: 51 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -522,7 +522,12 @@ def replicate_key(
522522

523523
self.update_primary_key_with_replica_keys(primary_key, replica_key, replica_region)
524524

525-
return ReplicateKeyResponse(ReplicaKeyMetadata=replica_key.metadata)
525+
# CurrentKeyMaterialId is not returned in the ReplicaKeyMetadata. May be due to not being evaluated until
526+
# the key has been successfully replicated as it does not show up in DescribeKey immediately either.
527+
replica_key_metadata_response = copy.deepcopy(replica_key.metadata)
528+
replica_key_metadata_response.pop("CurrentKeyMaterialId", None)
529+
530+
return ReplicateKeyResponse(ReplicaKeyMetadata=replica_key_metadata_response)
526531

527532
@staticmethod
528533
# Adds new multi region replica key to the primary key's metadata.
@@ -1206,13 +1211,10 @@ def import_key_material(
12061211
# TODO check if there was already a key imported for this kms key
12071212
# if so, it has to be identical. We cannot change keys by reimporting after deletion/expiry
12081213
key_material = self._decrypt_wrapped_key_material(import_state, encrypted_key_material)
1209-
1210-
if expiration_model:
1211-
key_to_import_material_to.metadata["ExpirationModel"] = expiration_model
1212-
else:
1213-
key_to_import_material_to.metadata["ExpirationModel"] = (
1214-
ExpirationModelType.KEY_MATERIAL_EXPIRES
1215-
)
1214+
key_material_id = key_to_import_material_to.generate_key_material_id(key_material)
1215+
key_to_import_material_to.metadata["ExpirationModel"] = (
1216+
expiration_model or ExpirationModelType.KEY_MATERIAL_EXPIRES
1217+
)
12161218
if (
12171219
key_to_import_material_to.metadata["ExpirationModel"]
12181220
== ExpirationModelType.KEY_MATERIAL_EXPIRES
@@ -1221,12 +1223,42 @@ def import_key_material(
12211223
raise ValidationException(
12221224
"A validTo date must be set if the ExpirationModel is KEY_MATERIAL_EXPIRES"
12231225
)
1226+
if existing_pending_material := key_to_import_material_to.crypto_key.pending_key_material:
1227+
pending_key_material_id = key_to_import_material_to.generate_key_material_id(
1228+
existing_pending_material
1229+
)
1230+
raise KMSInvalidStateException(
1231+
f"New key material (id: {key_material_id}) cannot be imported into KMS key "
1232+
f"{key_to_import_material_to.metadata['Arn']}, because another key material "
1233+
f"(id: {pending_key_material_id}) is pending rotation."
1234+
)
1235+
12241236
# TODO actually set validTo and make the key expire
12251237
key_to_import_material_to.metadata["Enabled"] = True
12261238
key_to_import_material_to.metadata["KeyState"] = KeyState.Enabled
12271239
key_to_import_material_to.crypto_key.load_key_material(key_material)
12281240

1229-
return ImportKeyMaterialResponse()
1241+
# KeyMaterialId / CurrentKeyMaterialId is only exposed for symmetric encryption keys.
1242+
key_material_id_response = None
1243+
if key_to_import_material_to.metadata["KeySpec"] == KeySpec.SYMMETRIC_DEFAULT:
1244+
key_material_id_response = key_to_import_material_to.generate_key_material_id(
1245+
key_material
1246+
)
1247+
1248+
# If there is no CurrentKeyMaterialId, instantly promote the pending key material to the current.
1249+
if key_to_import_material_to.metadata.get("CurrentKeyMaterialId") is None:
1250+
key_to_import_material_to.metadata["CurrentKeyMaterialId"] = (
1251+
key_material_id_response
1252+
)
1253+
key_to_import_material_to.crypto_key.key_material = (
1254+
key_to_import_material_to.crypto_key.pending_key_material
1255+
)
1256+
key_to_import_material_to.crypto_key.pending_key_material = None
1257+
1258+
return ImportKeyMaterialResponse(
1259+
KeyId=key_to_import_material_to.metadata["Arn"],
1260+
KeyMaterialId=key_material_id_response,
1261+
)
12301262

12311263
def delete_imported_key_material(
12321264
self,
@@ -1353,7 +1385,7 @@ def get_key_rotation_status(
13531385
key = self._get_kms_key(account_id, region_name, key_id, any_key_state_allowed=True)
13541386

13551387
response = GetKeyRotationStatusResponse(
1356-
KeyId=key_id,
1388+
KeyId=key.metadata["Arn"],
13571389
KeyRotationEnabled=key.is_key_rotation_enabled,
13581390
NextRotationDate=key.next_rotation_date,
13591391
)
@@ -1445,13 +1477,13 @@ def rotate_key_on_demand(
14451477

14461478
if key.metadata["KeySpec"] != KeySpec.SYMMETRIC_DEFAULT:
14471479
raise UnsupportedOperationException()
1448-
if key.metadata["Origin"] == OriginType.EXTERNAL:
1449-
raise NotImplementedError("Rotation of imported keys is not supported yet.")
1480+
self._validate_key_state_not_pending_import(key)
1481+
self._validate_external_key_has_pending_material(key)
14501482

14511483
key.rotate_key_on_demand()
14521484

14531485
return RotateKeyOnDemandResponse(
1454-
KeyId=key_id,
1486+
KeyId=key.metadata["Arn"],
14551487
)
14561488

14571489
@handler("TagResource", expand=False)
@@ -1528,6 +1560,12 @@ def _validate_key_state_not_pending_import(self, key: KmsKey):
15281560
if key.metadata["KeyState"] == KeyState.PendingImport:
15291561
raise KMSInvalidStateException(f"{key.metadata['Arn']} is pending import.")
15301562

1563+
def _validate_external_key_has_pending_material(self, key: KmsKey):
1564+
if key.metadata["Origin"] == "EXTERNAL" and key.crypto_key.pending_key_material is None:
1565+
raise KMSInvalidStateException(
1566+
f"No available key material pending rotation for the key: {key.metadata['Arn']}."
1567+
)
1568+
15311569
def _validate_key_for_encryption_decryption(self, context: RequestContext, key: KmsKey):
15321570
key_usage = key.metadata["KeyUsage"]
15331571
if key_usage != "ENCRYPT_DECRYPT":

‎localstack-core/localstack/testing/snapshots/transformer_utility.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -566,6 +566,8 @@ def kms_api():
566566
"""
567567
return [
568568
TransformerUtility.key_value("KeyId"),
569+
TransformerUtility.key_value("KeyMaterialId"),
570+
TransformerUtility.key_value("CurrentKeyMaterialId"),
569571
TransformerUtility.jsonpath(
570572
jsonpath="$..Signature",
571573
value_replacement="<signature>",

0 commit comments

Comments
 (0)