diff --git a/docs/docs/services/kms.rst b/docs/docs/services/kms.rst index f05f878016cb..765120c025d5 100644 --- a/docs/docs/services/kms.rst +++ b/docs/docs/services/kms.rst @@ -34,7 +34,7 @@ These authorization checks are quite basic for now. Moto will only throw an Acce Delete the alias. - [ ] delete_custom_key_store -- [ ] delete_imported_key_material +- [X] delete_imported_key_material - [ ] derive_shared_secret - [ ] describe_custom_key_stores - [X] describe_key @@ -53,9 +53,12 @@ These authorization checks are quite basic for now. Moto will only throw an Acce - [ ] get_key_last_usage - [X] get_key_policy - [X] get_key_rotation_status -- [ ] get_parameters_for_import +- [X] get_parameters_for_import + + Supported wrapping algorithms: RSAES_OAEP_SHA_256, RSAES_OAEP_SHA_1. + RSA_AES_KEY_WRAP variants are not yet implemented. - [X] get_public_key -- [ ] import_key_material +- [X] import_key_material - [X] list_aliases - [X] list_grants - [X] list_key_policies diff --git a/moto/kms/exceptions.py b/moto/kms/exceptions.py index 5a8ad4fc3e23..9c0f07f80957 100644 --- a/moto/kms/exceptions.py +++ b/moto/kms/exceptions.py @@ -65,3 +65,24 @@ def __init__(self) -> None: super().__init__("KMSInvalidMacException", "") self.description = '{"__type":"KMSInvalidMacException"}' + + +class UnsupportedOperationException(JsonRESTError): + code = 400 + + def __init__(self, message: str): + super().__init__("UnsupportedOperationException", message) + + +class KMSInvalidStateException(JsonRESTError): + code = 400 + + def __init__(self, message: str): + super().__init__("KMSInvalidStateException", message) + + +class InvalidImportTokenException(JsonRESTError): + code = 400 + + def __init__(self, message: str): + super().__init__("InvalidImportTokenException", message) diff --git a/moto/kms/models.py b/moto/kms/models.py index 585f6b8d52e9..a11682f8976c 100644 --- a/moto/kms/models.py +++ b/moto/kms/models.py @@ -17,13 +17,18 @@ from .exceptions import ( AccessDeniedException, + InvalidCiphertextException, + InvalidImportTokenException, InvalidKeyUsageException, KMSInvalidMacException, + KMSInvalidStateException, + UnsupportedOperationException, ValidationException, ) from .utils import ( RESERVED_ALIASES, KeySpec, + RSAWrappingKey, SigningAlgorithm, decrypt, encrypt, @@ -159,8 +164,13 @@ def __init__( } self.key_rotation_status = False self.deletion_date: datetime | None = None - self.key_material = generate_master_key() self.origin = origin + if self.origin == "EXTERNAL": + self.key_material: bytes | None = None + self.key_state = "PendingImport" + self.enabled = False + else: + self.key_material = generate_master_key() self.key_manager = "CUSTOMER" self.key_spec = key_spec or "SYMMETRIC_DEFAULT" self.private_key = generate_private_key(self.key_spec) @@ -172,6 +182,11 @@ def __init__( self.rotations: list[dict[str, Any]] = [] self.aliases: dict[str, Alias] = {} + # Import key material fields + self.import_token: bytes | None = None + self.wrapping_private_key: RSAWrappingKey | None = None + self.wrapping_algorithm: str | None = None + def add_grant( self, name: str, @@ -590,6 +605,7 @@ def encrypt( self, key_id: str, plaintext: bytes, encryption_context: dict[str, str] ) -> tuple[bytes, str]: key_id = self.any_id_to_key_id(key_id) + key = self.keys[key_id] ciphertext_blob = encrypt( master_keys=self.keys, @@ -597,7 +613,7 @@ def encrypt( plaintext=plaintext, encryption_context=encryption_context, ) - arn = self.keys[key_id].arn + arn = key.arn return ciphertext_blob, arn def decrypt( @@ -677,6 +693,124 @@ def generate_data_key_without_plaintext(self) -> None: # Responses uses 'generate_data_key' pass + def get_parameters_for_import( + self, key_id: str, wrapping_algorithm: str, wrapping_key_spec: str + ) -> tuple[bytes, bytes, float]: + """ + Supported wrapping algorithms: RSAES_OAEP_SHA_256, RSAES_OAEP_SHA_1. + RSA_AES_KEY_WRAP variants are not yet implemented. + """ + key_id = self.any_id_to_key_id(key_id) + key = self.keys[key_id] + + if key.origin != "EXTERNAL": + raise UnsupportedOperationException( + "The request was rejected because the specified KMS key cannot " + "accept imported key material. The Origin of the KMS key must be EXTERNAL." + ) + + if key.key_state not in ("PendingImport", "Enabled"): + raise KMSInvalidStateException( + f"arn:aws:kms:{key.region}:{key.account_id}:key/{key.id} is pending deletion." + ) + + # Validate wrapping key spec and generate wrapping key + if wrapping_key_spec == "RSA_2048": + key_size = 2048 + elif wrapping_key_spec == "RSA_3072": + key_size = 3072 + elif wrapping_key_spec == "RSA_4096": + key_size = 4096 + else: + raise ValidationException( + f"1 validation error detected: Value '{wrapping_key_spec}' at 'wrappingKeySpec' " + "failed to satisfy constraint: Member must satisfy enum value set: " + "[RSA_2048, RSA_3072, RSA_4096]" + ) + + wrapping_key = RSAWrappingKey(key_size) + + # Store the wrapping key and algorithm on the key for later use + key.wrapping_private_key = wrapping_key + key.wrapping_algorithm = wrapping_algorithm + + # Generate import token + key.import_token = os.urandom(32) + + # Expiration: 24 hours from now + parameters_valid_to = unix_time(utcnow() + timedelta(days=1)) + + return wrapping_key.public_key(), key.import_token, parameters_valid_to + + def import_key_material( + self, + key_id: str, + import_token: bytes, + encrypted_key_material: bytes, + expiration_model: str, + valid_to: float | None, + ) -> None: + key_id = self.any_id_to_key_id(key_id) + key = self.keys[key_id] + + if key.origin != "EXTERNAL": + raise UnsupportedOperationException( + "The request was rejected because the specified KMS key cannot " + "accept imported key material. The Origin of the KMS key must be EXTERNAL." + ) + + if key.key_state not in ("PendingImport", "Enabled"): + raise KMSInvalidStateException( + f"arn:aws:kms:{key.region}:{key.account_id}:key/{key.id} is not in a valid " + "state for this operation." + ) + + # Validate import token + if key.import_token is None or import_token != key.import_token: + raise InvalidImportTokenException( + "The request was rejected because the provided import token is " + "invalid or is associated with a different KMS key." + ) + + # Validate wrapping key exists + if key.wrapping_private_key is None: + raise InvalidImportTokenException( + "The request was rejected because the provided import token is " + "invalid or is associated with a different KMS key." + ) + + # Decrypt the encrypted key material using the stored wrapping key + try: + plaintext_key_material = key.wrapping_private_key.unwrap( + encrypted_key_material, key.wrapping_algorithm + ) + except Exception: + raise InvalidCiphertextException() + + # Set the key material + key.key_material = plaintext_key_material + key.key_state = "Enabled" + key.enabled = True + + def delete_imported_key_material(self, key_id: str) -> None: + key_id = self.any_id_to_key_id(key_id) + key = self.keys[key_id] + + if key.origin != "EXTERNAL": + raise UnsupportedOperationException( + "The request was rejected because the specified KMS key cannot " + "have its imported key material deleted. The Origin of the KMS key must be EXTERNAL." + ) + + if key.key_state in ("PendingDeletion",): + raise KMSInvalidStateException( + f"arn:aws:kms:{key.region}:{key.account_id}:key/{key.id} is pending deletion." + ) + + key.key_material = None + key.key_state = "PendingImport" + key.enabled = False + def list_resource_tags(self, key_id_or_arn: str) -> dict[str, list[dict[str, str]]]: key_id = self.get_key_id(key_id_or_arn) if key_id in self.keys: diff --git a/moto/kms/responses.py b/moto/kms/responses.py index 3edadeef36ac..f32224e0fc58 100644 --- a/moto/kms/responses.py +++ b/moto/kms/responses.py @@ -25,7 +25,13 @@ def __init__(self) -> None: def _get_param(self, param_name: str, if_none: Any = None) -> Any: params = json.loads(self.body) - for key in ("Plaintext", "CiphertextBlob", "Message"): + for key in ( + "Plaintext", + "CiphertextBlob", + "Message", + "EncryptedKeyMaterial", + "ImportToken", + ): if key in params: params[key] = base64.b64decode(params[key].encode("utf-8")) @@ -775,6 +781,63 @@ def get_public_key(self) -> str: } ) + def get_parameters_for_import(self) -> str: + """https://docs.aws.amazon.com/kms/latest/APIReference/API_GetParametersForImport.html""" + key_id = self._get_param("KeyId") + wrapping_algorithm = self._get_param("WrappingAlgorithm") + wrapping_key_spec = self._get_param("WrappingKeySpec") + + self._validate_key_id(key_id) + + public_key, import_token, parameters_valid_to = ( + self.kms_backend.get_parameters_for_import( + key_id=key_id, + wrapping_algorithm=wrapping_algorithm, + wrapping_key_spec=wrapping_key_spec, + ) + ) + + return json.dumps( + { + "KeyId": key_id, + "ImportToken": base64.b64encode(import_token).decode("utf-8"), + "PublicKey": base64.b64encode(public_key).decode("utf-8"), + "ParametersValidTo": parameters_valid_to, + } + ) + + def import_key_material(self) -> str: + """https://docs.aws.amazon.com/kms/latest/APIReference/API_ImportKeyMaterial.html""" + key_id = self._get_param("KeyId") + import_token = self._get_param("ImportToken") + encrypted_key_material = self._get_param("EncryptedKeyMaterial") + expiration_model = self._get_param( + "ExpirationModel", "KEY_MATERIAL_DOES_NOT_EXPIRE" + ) + valid_to = self._get_param("ValidTo") + + self._validate_key_id(key_id) + + self.kms_backend.import_key_material( + key_id=key_id, + import_token=import_token, + encrypted_key_material=encrypted_key_material, + expiration_model=expiration_model, + valid_to=valid_to, + ) + + return json.dumps({"KeyId": key_id}) + + def delete_imported_key_material(self) -> str: + """https://docs.aws.amazon.com/kms/latest/APIReference/API_DeleteImportedKeyMaterial.html""" + key_id = self._get_param("KeyId") + + self._validate_key_id(key_id) + + self.kms_backend.delete_imported_key_material(key_id=key_id) + + return "{}" + def rotate_key_on_demand(self) -> str: key_id = self._get_param("KeyId") diff --git a/moto/kms/utils.py b/moto/kms/utils.py index 881d17bbb94f..990d76b7b038 100644 --- a/moto/kms/utils.py +++ b/moto/kms/utils.py @@ -21,6 +21,7 @@ from .exceptions import ( AccessDeniedException, InvalidCiphertextException, + KMSInvalidStateException, NotFoundException, ValidationException, ) @@ -243,6 +244,51 @@ def public_key(self) -> bytes: ) +class RSAWrappingKey: + """RSA key used for wrapping (encrypting) key material during KMS import.""" + + __supported_key_sizes = [2048, 3072, 4096] + + def __init__(self, key_size: int): + if key_size not in self.__supported_key_sizes: + raise ValidationException( + f"1 validation error detected: Value '{key_size}' at 'wrappingKeySpec' " + "failed to satisfy constraint: Member must satisfy enum value set: " + f"{self.__supported_key_sizes}" + ) + self.key_size = key_size + self.private_key = rsa.generate_private_key( + public_exponent=65537, key_size=self.key_size + ) + + def public_key(self) -> bytes: + return self.private_key.public_key().public_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ) + + def unwrap(self, encrypted_material: bytes, wrapping_algorithm: str) -> bytes: + if wrapping_algorithm == "RSAES_OAEP_SHA_256": + pad = padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ) + elif wrapping_algorithm == "RSAES_OAEP_SHA_1": + pad = padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA1()), + algorithm=hashes.SHA1(), + label=None, + ) + else: + raise ValidationException( + f"1 validation error detected: Value '{wrapping_algorithm}' at 'wrappingAlgorithm' " + "failed to satisfy constraint: Member must satisfy enum value set: " + "[RSAES_OAEP_SHA_256, RSAES_OAEP_SHA_1]" + ) + return self.private_key.decrypt(encrypted_material, pad) + + class ECDSAPrivateKey(AbstractPrivateKey): def __init__(self, key_spec: str): validate_key_spec(key_spec, KeySpec.ecc_key_specs()) @@ -372,6 +418,11 @@ def encrypt( id_type = "Alias" if is_alias else "keyId" raise NotFoundException(f"{id_type} {key_id} is not found.") + if key.key_material is None: + raise KMSInvalidStateException( + f"{key_id} is not in a valid state for this operation." + ) + if plaintext == b"": raise ValidationException( "1 validation error detected: Value at 'plaintext' failed to satisfy constraint: Member must have length greater than or equal to 1" @@ -428,6 +479,11 @@ def decrypt( "does not exist in this region, or you are not allowed to access." ) + if key.key_material is None: + raise KMSInvalidStateException( + f"{ciphertext.key_id} is not in a valid state for this operation." + ) + try: decryptor = Cipher( algorithms.AES(key.key_material), diff --git a/tests/test_kms/test_kms_import_key_material.py b/tests/test_kms/test_kms_import_key_material.py new file mode 100644 index 000000000000..edbcffaaa735 --- /dev/null +++ b/tests/test_kms/test_kms_import_key_material.py @@ -0,0 +1,316 @@ +import boto3 +import pytest +from botocore.exceptions import ClientError +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import padding +from cryptography.hazmat.primitives.serialization import load_der_public_key + +from moto import mock_aws + +STATIC_KEY_MATERIAL = b"\x00" * 32 # 256-bit key for testing + + +def _encrypt_key_material(public_key_bytes: bytes, key_material: bytes) -> bytes: + """Encrypt key material using the wrapping public key with OAEP SHA-256.""" + public_key = load_der_public_key(public_key_bytes) + return public_key.encrypt( + key_material, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + + +@mock_aws +def test_create_key_with_external_origin(): + """A key with EXTERNAL origin should be in PendingImport state with no key material.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + + assert key["Origin"] == "EXTERNAL" + assert key["KeyState"] == "PendingImport" + assert key["Enabled"] is False + + +@mock_aws +def test_get_parameters_for_import_happy_path(): + """get_parameters_for_import returns a public key and import token.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + key_id = key["KeyId"] + + response = client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + + assert "PublicKey" in response + assert "ImportToken" in response + assert "ParametersValidTo" in response + assert response["KeyId"] == key_id + + # Verify public key is valid DER + public_key = load_der_public_key(response["PublicKey"]) + assert public_key.key_size == 2048 + + +@mock_aws +def test_get_parameters_for_import_non_external_key(): + """get_parameters_for_import should fail on a key with AWS_KMS origin.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key()["KeyMetadata"] + key_id = key["KeyId"] + + with pytest.raises(ClientError) as exc: + client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + + err = exc.value.response["Error"] + assert err["Code"] == "UnsupportedOperationException" + + +@mock_aws +def test_import_key_material_happy_path(): + """Full flow: create EXTERNAL key, get params, import material, encrypt/decrypt.""" + client = boto3.client("kms", region_name="us-east-1") + + # Create EXTERNAL key + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + key_id = key["KeyId"] + + # Get wrapping parameters + params = client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + + # Encrypt key material with the public key + encrypted_key_material = _encrypt_key_material( + params["PublicKey"], STATIC_KEY_MATERIAL + ) + + # Import key material + client.import_key_material( + KeyId=key_id, + ImportToken=params["ImportToken"], + EncryptedKeyMaterial=encrypted_key_material, + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + # Verify key is now enabled + key_desc = client.describe_key(KeyId=key_id)["KeyMetadata"] + assert key_desc["KeyState"] == "Enabled" + assert key_desc["Enabled"] is True + + # Verify encrypt/decrypt works + plaintext = b"Hello, World!" + encrypt_response = client.encrypt(KeyId=key_id, Plaintext=plaintext) + decrypt_response = client.decrypt(CiphertextBlob=encrypt_response["CiphertextBlob"]) + assert decrypt_response["Plaintext"] == plaintext + + +@mock_aws +def test_import_key_material_non_external_key(): + """import_key_material should fail on a key with AWS_KMS origin.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key()["KeyMetadata"] + key_id = key["KeyId"] + + with pytest.raises(ClientError) as exc: + client.import_key_material( + KeyId=key_id, + ImportToken=b"fake-token", + EncryptedKeyMaterial=b"fake-material", + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + err = exc.value.response["Error"] + assert err["Code"] == "UnsupportedOperationException" + + +@mock_aws +def test_import_key_material_invalid_token(): + """import_key_material should fail with a wrong import token.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + key_id = key["KeyId"] + + # Get params to initialize wrapping key + client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + + with pytest.raises(ClientError) as exc: + client.import_key_material( + KeyId=key_id, + ImportToken=b"wrong-token", + EncryptedKeyMaterial=b"fake-material", + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + err = exc.value.response["Error"] + assert err["Code"] == "InvalidImportTokenException" + + +@mock_aws +def test_delete_imported_key_material(): + """delete_imported_key_material should reset key to PendingImport state.""" + client = boto3.client("kms", region_name="us-east-1") + + # Create and import key material + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + key_id = key["KeyId"] + + params = client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + + encrypted_key_material = _encrypt_key_material( + params["PublicKey"], STATIC_KEY_MATERIAL + ) + + client.import_key_material( + KeyId=key_id, + ImportToken=params["ImportToken"], + EncryptedKeyMaterial=encrypted_key_material, + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + # Verify key is enabled + key_desc = client.describe_key(KeyId=key_id)["KeyMetadata"] + assert key_desc["KeyState"] == "Enabled" + + # Delete imported key material + client.delete_imported_key_material(KeyId=key_id) + + # Verify key is back to PendingImport + key_desc = client.describe_key(KeyId=key_id)["KeyMetadata"] + assert key_desc["KeyState"] == "PendingImport" + assert key_desc["Enabled"] is False + + # Verify encrypt no longer works + with pytest.raises(ClientError): + client.encrypt(KeyId=key_id, Plaintext=b"test") + # Key in PendingImport state should not be usable + + +@mock_aws +def test_delete_imported_key_material_non_external_key(): + """delete_imported_key_material should fail on a key with AWS_KMS origin.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key()["KeyMetadata"] + key_id = key["KeyId"] + + with pytest.raises(ClientError) as exc: + client.delete_imported_key_material(KeyId=key_id) + + err = exc.value.response["Error"] + assert err["Code"] == "UnsupportedOperationException" + + +@mock_aws +def test_reimport_same_key_material(): + """Reimporting key material into an already-enabled key should succeed.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + key_id = key["KeyId"] + + # First import + params = client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + encrypted_key_material = _encrypt_key_material( + params["PublicKey"], STATIC_KEY_MATERIAL + ) + client.import_key_material( + KeyId=key_id, + ImportToken=params["ImportToken"], + EncryptedKeyMaterial=encrypted_key_material, + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + # Encrypt something + plaintext = b"test data" + encrypt_response = client.encrypt(KeyId=key_id, Plaintext=plaintext) + + # Reimport the same key material (need fresh params) + params2 = client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_256", + WrappingKeySpec="RSA_2048", + ) + encrypted_key_material2 = _encrypt_key_material( + params2["PublicKey"], STATIC_KEY_MATERIAL + ) + client.import_key_material( + KeyId=key_id, + ImportToken=params2["ImportToken"], + EncryptedKeyMaterial=encrypted_key_material2, + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + # Should still be able to decrypt with the same key material + decrypt_response = client.decrypt(CiphertextBlob=encrypt_response["CiphertextBlob"]) + assert decrypt_response["Plaintext"] == plaintext + + +@mock_aws +def test_import_key_material_with_sha1_wrapping(): + """Import key material using RSAES_OAEP_SHA_1 wrapping algorithm.""" + client = boto3.client("kms", region_name="us-east-1") + + key = client.create_key(Origin="EXTERNAL")["KeyMetadata"] + key_id = key["KeyId"] + + params = client.get_parameters_for_import( + KeyId=key_id, + WrappingAlgorithm="RSAES_OAEP_SHA_1", + WrappingKeySpec="RSA_2048", + ) + + # Encrypt with SHA-1 OAEP + public_key = load_der_public_key(params["PublicKey"]) + encrypted_key_material = public_key.encrypt( + STATIC_KEY_MATERIAL, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA1()), + algorithm=hashes.SHA1(), + label=None, + ), + ) + + client.import_key_material( + KeyId=key_id, + ImportToken=params["ImportToken"], + EncryptedKeyMaterial=encrypted_key_material, + ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE", + ) + + # Verify works + key_desc = client.describe_key(KeyId=key_id)["KeyMetadata"] + assert key_desc["KeyState"] == "Enabled" + + plaintext = b"SHA1 wrapping test" + encrypt_response = client.encrypt(KeyId=key_id, Plaintext=plaintext) + decrypt_response = client.decrypt(CiphertextBlob=encrypt_response["CiphertextBlob"]) + assert decrypt_response["Plaintext"] == plaintext diff --git a/tests/test_kms/test_utils.py b/tests/test_kms/test_utils.py index 458abb53cbb6..c00c5a88709e 100644 --- a/tests/test_kms/test_utils.py +++ b/tests/test_kms/test_utils.py @@ -1,4 +1,7 @@ import pytest +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import padding +from cryptography.hazmat.primitives.serialization import load_der_public_key from moto.kms.exceptions import ( AccessDeniedException, @@ -13,6 +16,7 @@ ECDSAPrivateKey, KeySpec, RSAPrivateKey, + RSAWrappingKey, SigningAlgorithm, _deserialize_ciphertext_blob, _serialize_ciphertext_blob, @@ -243,3 +247,70 @@ def test_decrypt_invalid_encryption_context(): ciphertext_blob=ciphertext_blob, encryption_context={}, ) + + +# RSAWrappingKey tests + + +def test_rsa_wrapping_key_supports_all_key_sizes(): + for key_size in (2048, 3072, 4096): + wrapping_key = RSAWrappingKey(key_size) + public_key = load_der_public_key(wrapping_key.public_key()) + assert public_key.key_size == key_size + + +def test_rsa_wrapping_key_rejects_invalid_key_size(): + with pytest.raises(ValidationException): + RSAWrappingKey(1024) + + +def test_rsa_wrapping_key_unwrap_sha256(): + wrapping_key = RSAWrappingKey(2048) + plaintext = b"\x00" * 32 # 256-bit key material + + # Encrypt with the public key using OAEP SHA-256 + public_key = load_der_public_key(wrapping_key.public_key()) + encrypted = public_key.encrypt( + plaintext, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + + # Unwrap should recover the original plaintext + result = wrapping_key.unwrap(encrypted, "RSAES_OAEP_SHA_256") + assert result == plaintext + + +def test_rsa_wrapping_key_unwrap_sha1(): + wrapping_key = RSAWrappingKey(2048) + plaintext = b"\xab" * 32 + + public_key = load_der_public_key(wrapping_key.public_key()) + encrypted = public_key.encrypt( + plaintext, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA1()), + algorithm=hashes.SHA1(), + label=None, + ), + ) + + result = wrapping_key.unwrap(encrypted, "RSAES_OAEP_SHA_1") + assert result == plaintext + + +def test_rsa_wrapping_key_unwrap_rejects_invalid_algorithm(): + wrapping_key = RSAWrappingKey(2048) + + with pytest.raises(ValidationException): + wrapping_key.unwrap(b"fake_data", "RSA_AES_KEY_WRAP_SHA_256") + + +def test_rsa_wrapping_key_unwrap_rejects_bad_ciphertext(): + wrapping_key = RSAWrappingKey(2048) + + with pytest.raises(ValueError): + wrapping_key.unwrap(b"not_valid_ciphertext", "RSAES_OAEP_SHA_256")