fix

Author: RitvikKapila

examples/src/keyrings/aws_kms_discovery_multi_keyring_example.py

@@ -81,14 +81,16 @@ def get_aws_region_from_kms_key_id(kms_key_id: str) -> str:
 
 def encrypt_and_decrypt_with_keyring(
     kms_key_id: str,
-    aws_regions: list(str)
+    aws_regions: list[str]
 ):
     """Demonstrate an encrypt/decrypt cycle using an AWS KMS Discovery Multi Keyring.
 
-    Usage: encrypt_and_decrypt_with_keyring(kms_key_id)
+    Usage: encrypt_and_decrypt_with_keyring(kms_key_id, aws_regions)
     :param kms_key_id: KMS Key identifier for the KMS key you want to use for creating
     the kms_keyring used for encryption
     :type kms_key_id: string
+    :param aws_regions: List of AWS Regions to use for creating the discovery multi keyring
+    :type aws_regions: list[string]
 
     For more information on KMS Key identifiers, see
     https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id

examples/src/keyrings/aws_kms_multi_keyring_example.py

@@ -56,6 +56,19 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
+def get_aws_region_from_kms_key_id(kms_key_id: str) -> str:
+    """
+    Get the AWS Region from the KMS Key ID.
+
+    Usage: get_aws_region_from_kms_key_id(kms_key_id)
+    :param kms_key_id: KMS Key identifier for the KMS key you want to use
+    :type kms_key_id: string
+    :return: AWS Region
+    :rtype: string
+    """
+    return kms_key_id.split(":")[3]
+
+
 def encrypt_and_decrypt_with_keyring(
     default_region_kms_key_id: str,
     second_region_kms_key_id: str
@@ -144,7 +157,8 @@ def encrypt_and_decrypt_with_keyring(
     # `default_region_kms_key_id` directly.
 
     # 7a. Create a boto3 client for KMS for the default region.
-    default_region_kms_client = boto3.client('kms', region_name="us-west-2")
+    default_region = get_aws_region_from_kms_key_id(default_region_kms_key_id)
+    default_region_kms_client = boto3.client('kms', region_name=default_region)
 
     # 7b. Create KMS keyring
     default_region_kms_keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(
@@ -170,7 +184,8 @@ def encrypt_and_decrypt_with_keyring(
     # `second_region_kms_key_id` directly.
 
     # 8a. Create a boto3 client for KMS for the second region.
-    second_region_kms_client = boto3.client('kms', region_name="eu-central-1")
+    second_region = get_aws_region_from_kms_key_id(second_region_kms_key_id)
+    second_region_kms_client = boto3.client('kms', region_name=second_region)
 
     # 8b. Create KMS keyring
     second_region_kms_keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(

examples/src/keyrings/aws_kms_rsa_keyring_example.py

@@ -81,6 +81,9 @@ def encrypt_and_decrypt_with_keyring(
         config=MaterialProvidersConfig()
     )
 
+    # # Create the AWS KMS RSA keyring input
+    # For more information on the allowed encryption algorithms, please see
+    # https://github.com/awslabs/aws-encryption-sdk-specification/blob/master/framework/aws-kms/aws-kms-rsa-keyring.md#supported-padding-schemes
     keyring_input: CreateAwsKmsRsaKeyringInput = CreateAwsKmsRsaKeyringInput(
         public_key=public_key,
         kms_key_id=kms_rsa_key_id,

examples/test/keyrings/test_i_aws_kms_discovery_multi_keyring_example.py

@@ -11,4 +11,5 @@
 def test_encrypt_and_decrypt_with_keyring():
     """Test function for encrypt and decrypt using the AWS KMS Discovery Multi Keyring example."""
     kms_key_id = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
-    encrypt_and_decrypt_with_keyring(kms_key_id)
+    aws_regions = ["us-east-1", "us-west-2"]
+    encrypt_and_decrypt_with_keyring(kms_key_id, aws_regions)