From 78c518ddc832a30e1cf20015bc5c3b1850a1c797 Mon Sep 17 00:00:00 2001
From: Akira Noda <61897166+tsugumi-sys@users.noreply.github.com>
Date: Tue, 22 Aug 2023 03:07:12 +0900
Subject: [PATCH] KMS: Support 3072 and 4098 key sizes for RSA (#6708)

---
 moto/kms/utils.py                | 14 ++++++++++++++
 tests/test_kms/test_kms_boto3.py |  5 +++--
 tests/test_kms/test_utils.py     | 11 +++++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/moto/kms/utils.py b/moto/kms/utils.py
index 754eb7b2bc8d..89898d032e8d 100644
--- a/moto/kms/utils.py
+++ b/moto/kms/utils.py
@@ -165,7 +165,17 @@ def validate_signing_algorithm(
 
 
 class RSAPrivateKey(AbstractPrivateKey):
+    # See https://docs.aws.amazon.com/kms/latest/cryptographic-details/crypto-primitives.html
+    __supported_key_sizes = [2048, 3072, 4096]
+
     def __init__(self, key_size: int):
+        if key_size not in self.__supported_key_sizes:
+            raise ValidationException(
+                (
+                    "1 validation error detected: Value at 'key_size' failed "
+                    "to satisfy constraint: Member must satisfy enum value set: {supported_key_sizes}"
+                ).format(supported_key_sizes=self.__supported_key_sizes)
+            )
         self.key_size = key_size
         self.private_key = rsa.generate_private_key(
             public_exponent=65537, key_size=self.key_size
@@ -222,6 +232,10 @@ def generate_private_key(key_spec: str) -> AbstractPrivateKey:
     """Generate a private key to be used on asymmetric sign/verify."""
     if key_spec == KeySpec.RSA_2048:
         return RSAPrivateKey(key_size=2048)
+    elif key_spec == KeySpec.RSA_3072:
+        return RSAPrivateKey(key_size=3072)
+    elif key_spec == KeySpec.RSA_4096:
+        return RSAPrivateKey(key_size=4096)
     else:
         return RSAPrivateKey(key_size=2048)
 
diff --git a/tests/test_kms/test_kms_boto3.py b/tests/test_kms/test_kms_boto3.py
index f11816cc144b..3895df617d5e 100644
--- a/tests/test_kms/test_kms_boto3.py
+++ b/tests/test_kms/test_kms_boto3.py
@@ -1163,11 +1163,12 @@ def test_sign_and_verify_ignoring_grant_tokens():
 
 
 @mock_kms
-def test_sign_and_verify_digest_message_type_RSASSA_PSS_SHA_256():
+@pytest.mark.parametrize("key_spec", ["RSA_2048", "RSA_3072", "RSA_4096"])
+def test_sign_and_verify_digest_message_type_RSASSA_PSS_SHA_256(key_spec):
     client = boto3.client("kms", region_name="us-west-2")
 
     key = client.create_key(
-        Description="sign-key", KeyUsage="SIGN_VERIFY", KeySpec="RSA_2048"
+        Description="sign-key", KeyUsage="SIGN_VERIFY", KeySpec=key_spec
     )
     key_id = key["KeyMetadata"]["KeyId"]
 
diff --git a/tests/test_kms/test_utils.py b/tests/test_kms/test_utils.py
index 0645ec12dd58..039b7ab36648 100644
--- a/tests/test_kms/test_utils.py
+++ b/tests/test_kms/test_utils.py
@@ -4,6 +4,7 @@
     AccessDeniedException,
     InvalidCiphertextException,
     NotFoundException,
+    ValidationException,
 )
 from moto.kms.models import Key
 from moto.kms.utils import (
@@ -18,6 +19,7 @@
     Ciphertext,
     KeySpec,
     SigningAlgorithm,
+    RSAPrivateKey,
 )
 
 ENCRYPTION_CONTEXT_VECTORS = [
@@ -95,6 +97,15 @@ def test_SigningAlgorithm_Enum():
     )
 
 
+def test_RSAPrivateKey_invalid_key_size():
+    with pytest.raises(ValidationException) as ex:
+        _ = RSAPrivateKey(key_size=100)
+    assert (
+        ex.value.message
+        == "1 validation error detected: Value at 'key_size' failed to satisfy constraint: Member must satisfy enum value set: [2048, 3072, 4096]"
+    )
+
+
 def test_generate_data_key():
     test = generate_data_key(123)