From e03c0117722a857c3239be385a5fb4c5c02983a0 Mon Sep 17 00:00:00 2001 From: sannya-singal Date: Fri, 11 Oct 2024 12:22:49 +0530 Subject: [PATCH 1/4] kms: implement DeriveSharedSecret operation --- .../localstack/services/kms/models.py | 42 +++++++++++++- .../localstack/services/kms/provider.py | 49 ++++++++++++++++ tests/aws/services/kms/test_kms.py | 46 +++++++++++++++ tests/aws/services/kms/test_kms.snapshot.json | 57 +++++++++++++++++++ .../aws/services/kms/test_kms.validation.json | 3 + 5 files changed, 195 insertions(+), 2 deletions(-) diff --git a/localstack-core/localstack/services/kms/models.py b/localstack-core/localstack/services/kms/models.py index 62550d389ebf3..003058ce42e67 100644 --- a/localstack-core/localstack/services/kms/models.py +++ b/localstack-core/localstack/services/kms/models.py @@ -21,14 +21,18 @@ from cryptography.hazmat.primitives.asymmetric.padding import PSS, PKCS1v15 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey from cryptography.hazmat.primitives.asymmetric.utils import Prehashed +from cryptography.hazmat.primitives.kdf.hkdf import HKDF from localstack.aws.api.kms import ( CreateAliasRequest, CreateGrantRequest, CreateKeyRequest, EncryptionContextType, + InvalidKeyUsageException, KeyMetadata, + KeySpec, KeyState, + KeyUsageType, KMSInvalidMacException, KMSInvalidSignatureException, MacAlgorithmSpec, @@ -361,6 +365,27 @@ def verify( # AWS itself raises this exception without any additional message. raise KMSInvalidSignatureException() + def derive_shared_secret(self, public_key: bytes) -> bytes: + key_spec = self.metadata.get("KeySpec") + if key_spec in ["ECC_NIST_P256", "ECC_SECG_P256K1"]: + algorithm = hashes.SHA256() + elif key_spec == "ECC_NIST_P384": + algorithm = hashes.SHA384() + elif key_spec == "ECC_NIST_P521": + algorithm = hashes.SHA512() + else: + raise InvalidKeyUsageException( + f"{self.metadata['Arn']} key usage is {self.metadata['KeyUsage']} which is not valid for DeriveSharedSecret." + ) + + return HKDF( + algorithm=algorithm, + salt=None, + info=b"", + length=algorithm.digest_size, + backend=default_backend(), + ).derive(public_key) + # This method gets called when a key is replicated to another region. It's meant to populate the required metadata # fields in a new replica key. def replicate_metadata( @@ -616,14 +641,27 @@ def _get_key_usage(self, request_key_usage: str, key_spec: str) -> str: raise ValidationException( "You must specify a KeyUsage value for all KMS keys except for symmetric encryption keys." ) - elif request_key_usage != "GENERATE_VERIFY_MAC": + elif request_key_usage != KeyUsageType.GENERATE_VERIFY_MAC: raise ValidationException( f"1 validation error detected: Value '{request_key_usage}' at 'keyUsage' " f"failed to satisfy constraint: Member must satisfy enum value set: " f"[ENCRYPT_DECRYPT, SIGN_VERIFY, GENERATE_VERIFY_MAC]" ) else: - return "GENERATE_VERIFY_MAC" + return KeyUsageType.GENERATE_VERIFY_MAC + elif request_key_usage == KeyUsageType.KEY_AGREEMENT: + if key_spec not in [ + KeySpec.ECC_NIST_P256, + KeySpec.ECC_NIST_P384, + KeySpec.ECC_NIST_P521, + KeySpec.ECC_SECG_P256K1, + KeySpec.SM2, + ]: + raise ValidationException( + f"KeyUsage {request_key_usage} is not compatible with KeySpec {key_spec}" + ) + else: + return request_key_usage else: return request_key_usage or "ENCRYPT_DECRYPT" diff --git a/localstack-core/localstack/services/kms/provider.py b/localstack-core/localstack/services/kms/provider.py index 45452ebc935bf..3ccd54c359c30 100644 --- a/localstack-core/localstack/services/kms/provider.py +++ b/localstack-core/localstack/services/kms/provider.py @@ -25,6 +25,7 @@ DateType, DecryptResponse, DeleteAliasRequest, + DeriveSharedSecretResponse, DescribeKeyRequest, DescribeKeyResponse, DisabledException, @@ -59,9 +60,11 @@ InvalidCiphertextException, InvalidGrantIdException, InvalidKeyUsageException, + KeyAgreementAlgorithmSpec, KeyIdType, KeySpec, KeyState, + KeyUsageType, KmsApi, KMSInvalidStateException, LimitType, @@ -79,8 +82,10 @@ MultiRegionKey, NotFoundException, NullableBooleanType, + OriginType, PlaintextType, PrincipalIdType, + PublicKeyType, PutKeyPolicyRequest, RecipientInfo, ReEncryptResponse, @@ -1346,6 +1351,50 @@ def untag_resource(self, context: RequestContext, request: UntagResourceRequest) # AWS doesn't seem to mind removal of a non-existent tag, so we do not raise any exception. key.tags.pop(tag_key, None) + def derive_shared_secret( + self, + context: RequestContext, + key_id: KeyIdType, + key_agreement_algorithm: KeyAgreementAlgorithmSpec, + public_key: PublicKeyType, + grant_tokens: GrantTokenList = None, + dry_run: NullableBooleanType = None, + recipient: RecipientInfo = None, + **kwargs, + ) -> DeriveSharedSecretResponse: + key = self._get_kms_key( + context.account_id, + context.region, + key_id, + enabled_key_allowed=True, + disabled_key_allowed=True, + ) + key_usage = key.metadata.get("KeyUsage") + key_origin = key.metadata.get("Origin") + + if key_usage != KeyUsageType.KEY_AGREEMENT: + raise InvalidKeyUsageException( + f"{key.metadata['Arn']} key usage is {key_usage} which is not valid for {context.operation.name}." + ) + + if key_agreement_algorithm != KeyAgreementAlgorithmSpec.ECDH: + raise ValidationException( + f"1 validation error detected: Value '{key_agreement_algorithm}' at 'keyAgreementAlgorithm' " + f"failed to satisfy constraint: Member must satisfy enum value set: [ECDH]" + ) + + # TODO: Verify the actual error raised + if key_origin not in [OriginType.AWS_KMS, OriginType.EXTERNAL]: + raise ValueError(f"Key origin: {key_origin} is not valid for {context.operation.name}.") + + shared_secret = key.derive_shared_secret(public_key) + return DeriveSharedSecretResponse( + KeyId=key_id, + SharedSecret=shared_secret, + KeyAgreementAlgorithm=key_agreement_algorithm, + KeyOrigin=key_origin, + ) + def _validate_key_state_not_pending_import(self, key: KmsKey): if key.metadata["KeyState"] == KeyState.PendingImport: raise KMSInvalidStateException(f"{key.metadata['Arn']} is pending import.") diff --git a/tests/aws/services/kms/test_kms.py b/tests/aws/services/kms/test_kms.py index d7465e638512a..e6a2bf03f4fcc 100644 --- a/tests/aws/services/kms/test_kms.py +++ b/tests/aws/services/kms/test_kms.py @@ -1322,6 +1322,52 @@ def test_get_parameters_for_import(self, kms_create_key, snapshot, aws_client): ) snapshot.match("response-error", e.value.response) + @markers.aws.validated + def test_derive_shared_secret(self, kms_create_key, aws_client, snapshot): + snapshot.add_transformer( + snapshot.transform.key_value("SharedSecret", reference_replacement=False) + ) + + # Create two keys and derive the shared secret + key1 = kms_create_key(KeySpec="ECC_NIST_P256", KeyUsage="KEY_AGREEMENT") + + key2 = kms_create_key(KeySpec="ECC_NIST_P256", KeyUsage="KEY_AGREEMENT") + pub_key2 = aws_client.kms.get_public_key(KeyId=key2["KeyId"])["PublicKey"] + + secret = aws_client.kms.derive_shared_secret( + KeyId=key1["KeyId"], KeyAgreementAlgorithm="ECDH", PublicKey=pub_key2 + ) + + snapshot.match("response", secret) + + # Create a key with invalid key usage + key3 = kms_create_key(KeySpec="ECC_NIST_P256", KeyUsage="SIGN_VERIFY") + with pytest.raises(ClientError) as e: + aws_client.kms.derive_shared_secret( + KeyId=key3["KeyId"], KeyAgreementAlgorithm="ECDH", PublicKey=pub_key2 + ) + snapshot.match("response-invalid-key-usage", e.value.response) + + # Create a key with invalid key spec + with pytest.raises(ClientError) as e: + kms_create_key(KeySpec="RSA_2048", KeyUsage="KEY_AGREEMENT") + snapshot.match("response-invalid-key-spec", e.value.response) + + # Create a key with invalid key agreement algorithm + with pytest.raises(ClientError) as e: + aws_client.kms.derive_shared_secret( + KeyId=key1["KeyId"], KeyAgreementAlgorithm="INVALID", PublicKey=pub_key2 + ) + snapshot.match("response-invalid-key-agreement-algo", e.value.response) + + # Create a symmetric and try to derive the shared secret + key4 = kms_create_key() + with pytest.raises(ClientError) as e: + aws_client.kms.derive_shared_secret( + KeyId=key4["KeyId"], KeyAgreementAlgorithm="ECDH", PublicKey=pub_key2 + ) + snapshot.match("response-invalid-key", e.value.response) + class TestKMSMultiAccounts: @markers.aws.needs_fixing diff --git a/tests/aws/services/kms/test_kms.snapshot.json b/tests/aws/services/kms/test_kms.snapshot.json index 2f1bc07fda890..5ed002148e7e0 100644 --- a/tests/aws/services/kms/test_kms.snapshot.json +++ b/tests/aws/services/kms/test_kms.snapshot.json @@ -1726,5 +1726,62 @@ "Origin": "AWS_KMS" } } + }, + "tests/aws/services/kms/test_kms.py::TestKMS::test_derive_shared_secret": { + "recorded-date": "11-10-2024, 06:11:26", + "recorded-content": { + "response": { + "KeyAgreementAlgorithm": "ECDH", + "KeyId": "", + "KeyOrigin": "AWS_KMS", + "SharedSecret": "shared-secret", + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 200 + } + }, + "response-invalid-key-usage": { + "Error": { + "Code": "InvalidKeyUsageException", + "Message": " key usage is SIGN_VERIFY which is not valid for DeriveSharedSecret." + }, + "message": " key usage is SIGN_VERIFY which is not valid for DeriveSharedSecret.", + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 400 + } + }, + "response-invalid-key-spec": { + "Error": { + "Code": "ValidationException", + "Message": "KeyUsage KEY_AGREEMENT is not compatible with KeySpec RSA_2048" + }, + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 400 + } + }, + "response-invalid-key-agreement-algo": { + "Error": { + "Code": "ValidationException", + "Message": "1 validation error detected: Value 'INVALID' at 'keyAgreementAlgorithm' failed to satisfy constraint: Member must satisfy enum value set: [ECDH]" + }, + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 400 + } + }, + "response-invalid-key": { + "Error": { + "Code": "InvalidKeyUsageException", + "Message": " key usage is ENCRYPT_DECRYPT which is not valid for DeriveSharedSecret." + }, + "message": " key usage is ENCRYPT_DECRYPT which is not valid for DeriveSharedSecret.", + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 400 + } + } + } } } diff --git a/tests/aws/services/kms/test_kms.validation.json b/tests/aws/services/kms/test_kms.validation.json index 85260599d624c..7c47de2aa78d4 100644 --- a/tests/aws/services/kms/test_kms.validation.json +++ b/tests/aws/services/kms/test_kms.validation.json @@ -29,6 +29,9 @@ "tests/aws/services/kms/test_kms.py::TestKMS::test_create_multi_region_key": { "last_validated_date": "2024-04-11T15:53:40+00:00" }, + "tests/aws/services/kms/test_kms.py::TestKMS::test_derive_shared_secret": { + "last_validated_date": "2024-10-11T06:11:24+00:00" + }, "tests/aws/services/kms/test_kms.py::TestKMS::test_describe_and_list_sign_key": { "last_validated_date": "2024-04-11T15:53:27+00:00" }, From e937619a2dce7cc8e76d5b403aedf5c854b38fc5 Mon Sep 17 00:00:00 2001 From: sannya-singal Date: Fri, 11 Oct 2024 12:39:31 +0530 Subject: [PATCH 2/4] re-run the test --- tests/aws/services/kms/test_kms.snapshot.json | 2 +- tests/aws/services/kms/test_kms.validation.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/aws/services/kms/test_kms.snapshot.json b/tests/aws/services/kms/test_kms.snapshot.json index 5ed002148e7e0..0c3762fc42615 100644 --- a/tests/aws/services/kms/test_kms.snapshot.json +++ b/tests/aws/services/kms/test_kms.snapshot.json @@ -1728,7 +1728,7 @@ } }, "tests/aws/services/kms/test_kms.py::TestKMS::test_derive_shared_secret": { - "recorded-date": "11-10-2024, 06:11:26", + "recorded-date": "11-10-2024, 07:07:06", "recorded-content": { "response": { "KeyAgreementAlgorithm": "ECDH", diff --git a/tests/aws/services/kms/test_kms.validation.json b/tests/aws/services/kms/test_kms.validation.json index 7c47de2aa78d4..aa307db025382 100644 --- a/tests/aws/services/kms/test_kms.validation.json +++ b/tests/aws/services/kms/test_kms.validation.json @@ -30,7 +30,7 @@ "last_validated_date": "2024-04-11T15:53:40+00:00" }, "tests/aws/services/kms/test_kms.py::TestKMS::test_derive_shared_secret": { - "last_validated_date": "2024-10-11T06:11:24+00:00" + "last_validated_date": "2024-10-11T07:07:04+00:00" }, "tests/aws/services/kms/test_kms.py::TestKMS::test_describe_and_list_sign_key": { "last_validated_date": "2024-04-11T15:53:27+00:00" From 1a0eaeefeadca6b45b4ba27bf00c8651efb94628 Mon Sep 17 00:00:00 2001 From: sannya-singal Date: Fri, 11 Oct 2024 12:45:41 +0530 Subject: [PATCH 3/4] use specs from class --- localstack-core/localstack/services/kms/models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/localstack-core/localstack/services/kms/models.py b/localstack-core/localstack/services/kms/models.py index 003058ce42e67..f3e95a8826b40 100644 --- a/localstack-core/localstack/services/kms/models.py +++ b/localstack-core/localstack/services/kms/models.py @@ -367,11 +367,11 @@ def verify( def derive_shared_secret(self, public_key: bytes) -> bytes: key_spec = self.metadata.get("KeySpec") - if key_spec in ["ECC_NIST_P256", "ECC_SECG_P256K1"]: + if key_spec in [KeySpec.ECC_NIST_P256, KeySpec.ECC_SECG_P256K1]: algorithm = hashes.SHA256() - elif key_spec == "ECC_NIST_P384": + elif key_spec == KeySpec.ECC_NIST_P384: algorithm = hashes.SHA384() - elif key_spec == "ECC_NIST_P521": + elif key_spec == KeySpec.ECC_NIST_P521: algorithm = hashes.SHA512() else: raise InvalidKeyUsageException( From 3eba23655b3b7c4dafa1f0d04500ace739f3bb33 Mon Sep 17 00:00:00 2001 From: sannya-singal Date: Fri, 11 Oct 2024 15:40:54 +0530 Subject: [PATCH 4/4] use match case --- .../localstack/services/kms/models.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/localstack-core/localstack/services/kms/models.py b/localstack-core/localstack/services/kms/models.py index f3e95a8826b40..66decd56aad11 100644 --- a/localstack-core/localstack/services/kms/models.py +++ b/localstack-core/localstack/services/kms/models.py @@ -367,16 +367,17 @@ def verify( def derive_shared_secret(self, public_key: bytes) -> bytes: key_spec = self.metadata.get("KeySpec") - if key_spec in [KeySpec.ECC_NIST_P256, KeySpec.ECC_SECG_P256K1]: - algorithm = hashes.SHA256() - elif key_spec == KeySpec.ECC_NIST_P384: - algorithm = hashes.SHA384() - elif key_spec == KeySpec.ECC_NIST_P521: - algorithm = hashes.SHA512() - else: - raise InvalidKeyUsageException( - f"{self.metadata['Arn']} key usage is {self.metadata['KeyUsage']} which is not valid for DeriveSharedSecret." - ) + match key_spec: + case KeySpec.ECC_NIST_P256 | KeySpec.ECC_SECG_P256K1: + algorithm = hashes.SHA256() + case KeySpec.ECC_NIST_P384: + algorithm = hashes.SHA384() + case KeySpec.ECC_NIST_P521: + algorithm = hashes.SHA512() + case _: + raise InvalidKeyUsageException( + f"{self.metadata['Arn']} key usage is {self.metadata['KeyUsage']} which is not valid for DeriveSharedSecret." + ) return HKDF( algorithm=algorithm,