8000 kms: implement DeriveSharedSecret operation by sannya-singal · Pull Request #11672 · localstack/localstack · GitHub
[go: up one dir, main page]

Skip to content

kms: implement DeriveSharedSecret operation #11672

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Oct 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 41 additions & 2 deletions localstack-core/localstack/services/kms/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -361,6 +365,28 @@ 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")
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,
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(
Expand Down Expand Up @@ -616,14 +642,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"

Expand Down
49 changes: 49 additions & 0 deletions localstack-core/localstack/services/kms/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
DateType,
DecryptResponse,
DeleteAliasRequest,
DeriveSharedSecretResponse,
DescribeKeyRequest,
DescribeKeyResponse,
DisabledException,
Expand Down Expand Up @@ -59,9 +60,11 @@
InvalidCiphertextException,
InvalidGrantIdException,
InvalidKeyUsageException,
KeyAgreementAlgorithmSpec,
KeyIdType,
KeySpec,
KeyState,
KeyUsageType,
KmsApi,
KMSInvalidStateException,
LimitType,
Expand All @@ -79,8 +82,10 @@
MultiRegionKey,
NotFoundException,
NullableBooleanType,
OriginType,
PlaintextType,
PrincipalIdType,
PublicKeyType,
PutKeyPolicyRequest,
RecipientInfo,
ReEncryptResponse,
Expand Down Expand Up @@ -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.")
Expand Down
46 changes: 46 additions & 0 deletions tests/aws/services/kms/test_kms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
57 changes: 57 additions & 0 deletions tests/aws/services/kms/test_kms.snapshot.json
Original file line number Diff line number Diff line change
Expand Up @@ -1726,5 +1726,62 @@
"Origin": "AWS_KMS"
}
}
},
"tests/aws/services/kms/test_kms.py::TestKMS::test_derive_shared_secret": {
"recorded-date": "11-10-2024, 07:07:06",
"recorded-content": {
"response": {
"KeyAgreementAlgorithm": "ECDH",
"KeyId": "<key-id:1>",
"KeyOrigin": "AWS_KMS",
"SharedSecret": "shared-secret",
"ResponseMetadata": {
"HTTPHeaders": {},
"HTTPStatusCode": 200
}
},
"response-invalid-key-usage": {
"Error": {
"Code": "InvalidKeyUsageException",
"Message": "<key-arn> key usage is SIGN_VERIFY which is not valid for DeriveSharedSecret."
},
"message": "<key-arn> 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-arn> key usage is ENCRYPT_DECRYPT which is not valid for DeriveSharedSecret."
},
"message": "<key-arn> key usage is ENCRYPT_DECRYPT which is not valid for DeriveSharedSecret.",
"ResponseMetadata": {
"HTTPHeaders": {},
"HTTPStatusCode": 400
}
}
}
}
}
3 changes: 3 additions & 0 deletions tests/aws/services/kms/test_kms.validation.json
Original file line number Diff line number Diff line change
Expand Up @@ -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-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"
},
Expand Down
Loading
0