fix

Author: RitvikKapila

examples/src/keyrings/aws_kms_discovery_keyring_example.py

@@ -57,41 +57,21 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
-def get_account_id_from_kms_key_id(kms_key_id: str) -> str:
-    """
-    Get the AWS Account ID from the KMS Key ID.
-
-    Usage: get_account_id_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 Account ID
-    :rtype: string
-    """
-    return kms_key_id.split(":")[4]
-
-
-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(
-    kms_key_id: str
+    kms_key_id: str,
+    aws_account_id: str,
+    aws_region: str
 ):
     """Demonstrate an encrypt/decrypt cycle using an AWS KMS Discovery Keyring.
 
-    Usage: encrypt_and_decrypt_with_keyring(kms_key_id)
+    Usage: encrypt_and_decrypt_with_keyring(kms_key_id, aws_account_id)
     :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_account_id: AWS Account ID to use in the discovery filter
+    :type aws_account_id: string
+    :param aws_region: AWS Region to use for the kms client
+    :type aws_region: string
 
     For more information on KMS Key identifiers, see
     https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
@@ -108,11 +88,7 @@ def encrypt_and_decrypt_with_keyring(
     )
 
     # 2. Create a boto3 client for KMS.
-
-    # Get the AWS Region from the KMS Key ID.
-    region_name = get_aws_region_from_kms_key_id(kms_key_id)
-
-    kms_client = boto3.client('kms', region_name=region_name)
+    kms_client = boto3.client('kms', region_name=aws_region)
 
     # 3. Create encryption context.
     # Remember that your encryption context is NOT SECRET.
@@ -158,9 +134,6 @@ def encrypt_and_decrypt_with_keyring(
     # so that we limit the set of ciphertexts we are willing to decrypt to only ones
     # created by KMS keys in our account and partition.
 
-    # Get the AWS Account ID from the KMS Key ID to use in the discovery filter
-    aws_account_id: str = get_account_id_from_kms_key_id(kms_key_id=kms_key_id)
-
     discovery_keyring_input: CreateAwsKmsDiscoveryKeyringInput = CreateAwsKmsDiscoveryKeyringInput(
         kms_client=kms_client,
         discovery_filter=DiscoveryFilter(

examples/src/keyrings/aws_kms_discovery_multi_keyring_example.py

@@ -53,42 +53,19 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
-def get_account_id_from_kms_key_id(kms_key_id: str) -> str:
-    """
-    Get the AWS Account ID from the KMS Key ID.
-
-    Usage: get_account_id_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 Account ID
-    :rtype: string
-    """
-    return kms_key_id.split(":")[4]
-
-
-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(
     kms_key_id: str,
+    aws_account_id: 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, aws_regions)
+    Usage: encrypt_and_decrypt_with_keyring(kms_key_id, aws_account_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_account_id: AWS Account ID to use in the discovery filter
+    :type aws_account_id: string
     :param aws_regions: List of AWS Regions to use for creating the discovery multi keyring
     :type aws_regions: list[string]
 
@@ -107,11 +84,7 @@ def encrypt_and_decrypt_with_keyring(
     )
 
     # 2. Create a boto3 client for KMS.
-
-    # Get the AWS Region from the KMS Key ID.
-    region_name = get_aws_region_from_kms_key_id(kms_key_id)
-
-    kms_client = boto3.client('kms', region_name=region_name)
+    kms_client = boto3.client('kms', region_name="us-west-2")
 
     # 3. Create encryption context.
     # Remember that your encryption context is NOT SECRET.
@@ -156,10 +129,6 @@ def encrypt_and_decrypt_with_keyring(
     # 7. Now create a Discovery Multi keyring to use for decryption. We'll add a discovery filter
     # so that we limit the set of ciphertexts we are willing to decrypt to only ones
     # created by KMS keys in our account and partition.
-
-    # Get the AWS Account ID from the KMS Key ID to use in the discovery filter
-    aws_account_id: str = get_account_id_from_kms_key_id(kms_key_id=kms_key_id)
-
     discovery_multi_keyring_input: CreateAwsKmsDiscoveryMultiKeyringInput = \
         CreateAwsKmsDiscoveryMultiKeyringInput(
             regions=aws_regions,

examples/src/keyrings/aws_kms_multi_keyring_example.py

@@ -26,7 +26,7 @@
 with an encryption context. This example also includes some sanity checks for demonstration:
 1. Ciphertext and plaintext data are not the same
 2. Decryption of ciphertext is possible using the multi_keyring,
-    and every one of the keyrings from the multi_keyring separately
+   and every one of the keyrings from the multi_keyring separately
 3. All decrypted plaintext value match EXAMPLE_DATA
 These sanity checks are for demonstration in the example only. You do not need these in your code.
 
@@ -56,36 +56,32 @@
 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
+    second_region_kms_key_id: str,
+    default_region: str,
+    second_region: str
 ):
     """Demonstrate an encrypt/decrypt cycle using an AWS KMS Multi keyring.
     The multi_keyring is created using a KMS keyring as generator keyring and another KMS keyring
     as a child keyring. For this example, `default_region_kms_key_id` is the generator key id
     for a KMS key located in your default region, and `second_region_kms_key_id` is the KMS key id
-    for a KMS Key located in some second Region.
+    for a KMS Key located in some second region.
 
-    Usage: encrypt_and_decrypt_with_keyring(default_region_kms_key_id, second_region_kms_key_id)
+    Usage: encrypt_and_decrypt_with_keyring(default_region_kms_key_id,
+                                            second_region_kms_key_id,
+                                            default_region,
+                                            second_region)
     :param default_region_kms_key_id: KMS Key identifier for the default region KMS key you want to
     use as a generator keyring
     :type default_region_kms_key_id: string
     :param second_region_kms_key_id: KMS Key identifier for the second region KMS key you want to
     use as a child keyring
     :type second_region_kms_key_id: string
+    :param default_region: AWS Region for the default region KMS key
+    :type default_region: string
+    :param second_region: AWS Region for the second region KMS key
+    :type second_region: string
 
     For more information on KMS Key identifiers, see
     https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
@@ -155,9 +151,9 @@ def encrypt_and_decrypt_with_keyring(
 
     # 7. Demonstrate that you can successfully decrypt data using a KMS keyring with just the
     # `default_region_kms_key_id` directly.
+    # (This is an example for demonstration; you do not need to do this in your own code.)
 
     # 7a. Create a boto3 client for KMS for the default region.
-    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
@@ -182,9 +178,9 @@ def encrypt_and_decrypt_with_keyring(
 
     # 8. Demonstrate that you can also successfully decrypt data using a KMS keyring with just the
     # `second_region_kms_key_id` directly.
+    # (This is an example for demonstration; you do not need to do this in your own code.)
 
     # 8a. Create a boto3 client for KMS for the second region.
-    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

examples/src/keyrings/aws_kms_rsa_keyring_example.py

@@ -81,7 +81,7 @@ def encrypt_and_decrypt_with_keyring(
         config=MaterialProvidersConfig()
     )
 
-    # # Create the AWS KMS RSA keyring input
+    # Create the AWS KMS RSA keyring input
     # For more information on the allowed encryption algorithms, please see
     # https://docs.aws.amazon.com/kms/latest/developerguide/asymmetric-key-specs.html#key-spec-rsa
     keyring_input: CreateAwsKmsRsaKeyringInput = CreateAwsKmsRsaKeyringInput(

examples/src/keyrings/multi_keyring_example.py

@@ -182,6 +182,7 @@ def encrypt_and_decrypt_with_keyring(
 
     # 11. Demonstrate that you can successfully decrypt data using just the `kms_keyring`
     # directly.
+    # (This is an example for demonstration; you do not need to do this in your own code.)
 
     # 11a. Decrypt your encrypted data using the kms_keyring.
     plaintext_bytes_kms_keyring, _ = client.decrypt(
@@ -195,6 +196,7 @@ def encrypt_and_decrypt_with_keyring(
 
     # 12. Demonstrate that you can also successfully decrypt data using the `raw_aes_keyring`
     # directly.
+    # (This is an example for demonstration; you do not need to do this in your own code.)
 
     # 12a. Decrypt your encrypted data using the raw_aes_keyring.
     plaintext_bytes_raw_aes_keyring, _ = client.decrypt(

examples/test/keyrings/test_i_aws_kms_discovery_keyring_example.py

@@ -11,4 +11,6 @@
 def test_encrypt_and_decrypt_with_keyring():
     """Test function for encrypt and decrypt using the AWS KMS Discovery 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_account_id = "658956600833"
+    aws_region = "us-west-2"
+    encrypt_and_decrypt_with_keyring(kms_key_id, aws_account_id, aws_region)

examples/test/keyrings/test_i_aws_kms_discovery_multi_keyring_example.py

@@ -11,5 +11,6 @@
 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"
+    aws_account_id = "658956600833"
     aws_regions = ["us-east-1", "us-west-2"]
-    encrypt_and_decrypt_with_keyring(kms_key_id, aws_regions)
+    encrypt_and_decrypt_with_keyring(kms_key_id, aws_account_id, aws_regions)

examples/test/keyrings/test_i_aws_kms_multi_keyring_example.py

@@ -14,4 +14,9 @@ def test_encrypt_and_decrypt_with_keyring():
         "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
     second_region_kms_key_id = \
         "arn:aws:kms:eu-central-1:658956600833:key/75414c93-5285-4b57-99c9-30c1cf0a22c2"
-    encrypt_and_decrypt_with_keyring(default_region_kms_key_id, second_region_kms_key_id)
+    default_region = "us-west-2"
+    second_region = "eu-central-1"
+    encrypt_and_decrypt_with_keyring(default_region_kms_key_id,
+                                     second_region_kms_key_id,
+                                     default_region,
+                                     second_region)