test_kms.py

# Note/disclosure: This file has been (partially or fully) generated by an AI agent.
import boto3
import pytest
from botocore.exceptions import ClientError
from localstack.aws.connect import connect_to
from localstack.utils.strings import short_uid


def test_kms_key_operations(start_aws_proxy, cleanups):
    """Test basic KMS key operations with proxy."""
    key_description = f"test-key-{short_uid()}"

    # start proxy - forwarding requests for KMS keys with specific alias pattern
    config = {
        "services": {
            "kms": {"resources": [".*:key/.*"]},
        }
    }
    start_aws_proxy(config)

    # create clients
    kms_client = connect_to().kms
    kms_client_aws = boto3.client("kms")

    # create key in AWS
    create_response = kms_client_aws.create_key(Description=key_description)
    key_id_aws = create_response["KeyMetadata"]["KeyId"]
    key_arn_aws = create_response["KeyMetadata"]["Arn"]
    cleanups.append(
        lambda: kms_client_aws.schedule_key_deletion(
            KeyId=key_id_aws, PendingWindowInDays=7
        )
    )

    # assert that local call for this key is proxied
    key_aws = kms_client_aws.describe_key(KeyId=key_id_aws)
    key_local = kms_client.describe_key(KeyId=key_id_aws)
    assert key_local["KeyMetadata"]["KeyId"] == key_aws["KeyMetadata"]["KeyId"]
    assert key_local["KeyMetadata"]["Arn"] == key_arn_aws

    # test encryption with AWS client
    plaintext = b"test message"
    encrypt_response_aws = kms_client_aws.encrypt(KeyId=key_id_aws, Plaintext=plaintext)
    ciphertext = encrypt_response_aws["CiphertextBlob"]

    # decrypt with local client (proxied to AWS)
    decrypt_response_local = kms_client.decrypt(CiphertextBlob=ciphertext)
    assert decrypt_response_local["Plaintext"] == plaintext
    assert decrypt_response_local["KeyId"] == key_arn_aws

    # encrypt with local client (proxied to AWS)
    plaintext2 = b"another test message"
    encrypt_response_local = kms_client.encrypt(KeyId=key_id_aws, Plaintext=plaintext2)
    ciphertext2 = encrypt_response_local["CiphertextBlob"]

    # decrypt with AWS client
    decrypt_response_aws = kms_client_aws.decrypt(CiphertextBlob=ciphertext2)
    assert decrypt_response_aws["Plaintext"] == plaintext2


def test_kms_key_alias_operations(start_aws_proxy, cleanups):
    """Test KMS key alias operations with proxy."""
    alias_name = f"alias/test-alias-{short_uid()}"
    key_description = f"test-key-with-alias-{short_uid()}"

    # start proxy - forwarding requests for KMS
    config = {
        "services": {
            "kms": {"resources": [".*:key/.*", f".*:{alias_name}"]},
        }
    }
    start_aws_proxy(config)

    # create clients
    kms_client = connect_to().kms
    kms_client_aws = boto3.client("kms")

    # create key in AWS
    create_response = kms_client_aws.create_key(Description=key_description)
    key_id_aws = create_response["KeyMetadata"]["KeyId"]
    cleanups.append(
        lambda: kms_client_aws.schedule_key_deletion(
            KeyId=key_id_aws, PendingWindowInDays=7
        )
    )

    # create alias in AWS
    kms_client_aws.create_alias(AliasName=alias_name, TargetKeyId=key_id_aws)
    cleanups.append(lambda: kms_client_aws.delete_alias(AliasName=alias_name))

    # assert that local call for alias operations is proxied
    aliases_aws = kms_client_aws.list_aliases(KeyId=key_id_aws)["Aliases"]
    aliases_local = kms_client.list_aliases(KeyId=key_id_aws)["Aliases"]

    # filter for our specific alias
    alias_aws = [a for a in aliases_aws if a["AliasName"] == alias_name][0]
    alias_local = [a for a in aliases_local if a["AliasName"] == alias_name][0]

    assert alias_local["AliasName"] == alias_aws["AliasName"]
    assert alias_local["TargetKeyId"] == alias_aws["TargetKeyId"]

    # test encryption with alias via local client
    plaintext = b"test with alias"
    encrypt_response_local = kms_client.encrypt(KeyId=alias_name, Plaintext=plaintext)
    ciphertext = encrypt_response_local["CiphertextBlob"]

    # decrypt with AWS client
    decrypt_response_aws = kms_client_aws.decrypt(CiphertextBlob=ciphertext)
    assert decrypt_response_aws["Plaintext"] == plaintext


def test_kms_readonly_operations(start_aws_proxy, cleanups):
    """Test KMS operations in read-only proxy mode."""
    key_description = f"test-readonly-key-{short_uid()}"

    # start proxy - forwarding requests for KMS in read-only mode
    config = {
        "services": {
            "kms": {"resources": [".*:key/.*"], "read_only": True},
        }
    }
    start_aws_proxy(config)

    # create clients
    kms_client = connect_to().kms
    kms_client_aws = boto3.client("kms")

    # create key in AWS (this should succeed as it's direct AWS client)
    create_response = kms_client_aws.create_key(Description=key_description)
    key_id_aws = create_response["KeyMetadata"]["KeyId"]
    cleanups.append(
        lambda: kms_client_aws.schedule_key_deletion(
            KeyId=key_id_aws, PendingWindowInDays=7
        )
    )

    # assert that local call for describe_key is proxied and results are consistent
    key_aws = kms_client_aws.describe_key(KeyId=key_id_aws)
    key_local = kms_client.describe_key(KeyId=key_id_aws)
    assert key_local["KeyMetadata"]["KeyId"] == key_aws["KeyMetadata"]["KeyId"]

    # assert that local call for list_keys is proxied
    keys_local = kms_client.list_keys()["Keys"]
    keys_aws = kms_client_aws.list_keys()["Keys"]

    # filter for our specific key
    key_local_filtered = [k for k in keys_local if k["KeyId"] == key_id_aws]
    key_aws_filtered = [k for k in keys_aws if k["KeyId"] == key_id_aws]
    assert key_local_filtered == key_aws_filtered

    # Negative test: attempt write operations with proxied client
    # Create a new key using the proxied client (should succeed in LocalStack)
    new_key_description = f"no-proxy-key-{short_uid()}"
    new_key_local = kms_client.create_key(Description=new_key_description)
    new_key_id_local = new_key_local["KeyMetadata"]["KeyId"]
    cleanups.append(
        lambda: kms_client.schedule_key_deletion(
            KeyId=new_key_id_local, PendingWindowInDays=7
        )
    )

    # Verify that this new key does NOT exist in real AWS
    keys_aws_after_create = kms_client_aws.list_keys()["Keys"]
    assert not any(k for k in keys_aws_after_create if k["KeyId"] == new_key_id_local)

    # Attempt to encrypt data with the AWS key using proxied client (should fail)
    plaintext = b"this should not work"
    with pytest.raises(ClientError) as excinfo:
        kms_client.encrypt(KeyId=key_id_aws, Plaintext=plaintext)
    # In read-only mode, the key exists in AWS but LocalStack doesn't have it locally
    # so encrypt operation should fail
    assert excinfo.value.response["Error"]["Code"] in [
        "NotFoundException",
        "InvalidKeyId.NotFound",
    ]


def test_kms_selective_resource_matching(start_aws_proxy, cleanups):
    """Test that proxy forwards requests for specific KMS keys matching ARN pattern."""
    key_description_1 = f"test-proxied-key-1-{short_uid()}"
    key_description_2 = f"test-proxied-key-2-{short_uid()}"

    # create clients
    kms_client_aws = boto3.client("kms")

    # create two keys in AWS
    create_response_1 = kms_client_aws.create_key(Description=key_description_1)
    key_id_1 = create_response_1["KeyMetadata"]["KeyId"]
    key_arn_1 = create_response_1["KeyMetadata"]["Arn"]
    cleanups.append(
        lambda: kms_client_aws.schedule_key_deletion(
            KeyId=key_id_1, PendingWindowInDays=7
        )
    )

    create_response_2 = kms_client_aws.create_key(Description=key_description_2)
    key_id_2 = create_response_2["KeyMetadata"]["KeyId"]
    key_arn_2 = create_response_2["KeyMetadata"]["Arn"]
    cleanups.append(
        lambda: kms_client_aws.schedule_key_deletion(
            KeyId=key_id_2, PendingWindowInDays=7
        )
    )

    # start proxy - forwarding requests for both keys using wildcard pattern
    config = {
        "services": {
            "kms": {"resources": [".*:key/.*"]},
        }
    }
    start_aws_proxy(config)

    # create LocalStack client after proxy is started
    kms_client = connect_to().kms

    # verify that both keys are accessible via local client (proxied)
    key_1_local = kms_client.describe_key(KeyId=key_id_1)
    assert key_1_local["KeyMetadata"]["KeyId"] == key_id_1
    assert key_1_local["KeyMetadata"]["Arn"] == key_arn_1

    key_2_local = kms_client.describe_key(KeyId=key_id_2)
    assert key_2_local["KeyMetadata"]["KeyId"] == key_id_2
    assert key_2_local["KeyMetadata"]["Arn"] == key_arn_2

    # verify we can encrypt with both proxied keys
    plaintext = b"test with first key"
    encrypt_response_1 = kms_client.encrypt(KeyId=key_id_1, Plaintext=plaintext)
    ciphertext_1 = encrypt_response_1["CiphertextBlob"]

    # decrypt with AWS client to confirm it went through proxy
    decrypt_response_1 = kms_client_aws.decrypt(CiphertextBlob=ciphertext_1)
    assert decrypt_response_1["Plaintext"] == plaintext

    # encrypt with second key
    plaintext_2 = b"test with second key"
    encrypt_response_2 = kms_client.encrypt(KeyId=key_id_2, Plaintext=plaintext_2)
    ciphertext_2 = encrypt_response_2["CiphertextBlob"]

    # decrypt with AWS client
    decrypt_response_2 = kms_client_aws.decrypt(CiphertextBlob=ciphertext_2)
    assert decrypt_response_2["Plaintext"] == plaintext_2