updated ../../src/keyrings/raw_rsa_keyring_example.py

Author: RitvikKapila

examples/src/keyrings/raw_rsa_keyring_example.py

@@ -52,19 +52,42 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
-def generate_rsa_keyring():
-    """Creates a new RSA keyring along with generating new keys
+def encrypt_and_decrypt_with_keyring():
+    """Demonstrate an encrypt/decrypt cycle using a Raw RSA keyring.
 
-    Usage: generate_rsa_keyring()
+    Usage: encrypt_and_decrypt_with_keyring()
     """
-    # 1. The key namespace and key name are defined by you.
+    # 1. Instantiate the encryption SDK client.
+    # This builds the client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy,
+    # which enforces that this client only encrypts using committing algorithm suites and enforces
+    # that this client will only decrypt encrypted messages that were created with a committing
+    # algorithm suite.
+    # This is the default commitment policy if you were to build the client as
+    # `client = aws_encryption_sdk.EncryptionSDKClient()`.
+    client = aws_encryption_sdk.EncryptionSDKClient(
+        commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
+    )
+
+    # 2. Create encryption context.
+    # Remember that your encryption context is NOT SECRET.
+    # For more information, see
+    # https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
+    encryption_context: Dict[str, str] = {
+        "encryption": "context",
+        "is not": "secret",
+        "but adds": "useful metadata",
+        "that can help you": "be confident that",
+        "the data you are handling": "is what you think it is",
+    }
+
+    # 3. The key namespace and key name are defined by you.
     # and are used by the Raw RSA keyring
     # For more information, see
     # https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/use-raw-rsa-keyring.html
     key_name_space = "Some managed raw keys"
     key_name = "My 4096-bit RSA wrapping key"
 
-    # 2. Generate a 4096-bit RSA key to use with your keyring.
+    # 4. Generate a 4096-bit RSA key to use with your keyring.
     ssh_rsa_exponent = 65537
     bit_strength = 4096
     key = rsa.generate_private_key(
@@ -86,7 +109,7 @@ def generate_rsa_keyring():
         format=crypto_serialization.PublicFormat.SubjectPublicKeyInfo
     )
 
-    # 3. Create a Raw RSA keyring
+    # 5. Create a Raw RSA keyring
     mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
         config=MaterialProvidersConfig()
     )
@@ -103,64 +126,30 @@ def generate_rsa_keyring():
         input=keyring_input
     )
 
-    return raw_rsa_keyring
-
-
-def encrypt_and_decrypt_with_keyring():
-    """Demonstrate an encrypt/decrypt cycle using a Raw RSA keyring.
-
-    Usage: encrypt_and_decrypt_with_keyring()
-    """
-    # 1. Instantiate the encryption SDK client.
-    # This builds the client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy,
-    # which enforces that this client only encrypts using committing algorithm suites and enforces
-    # that this client will only decrypt encrypted messages that were created with a committing
-    # algorithm suite.
-    # This is the default commitment policy if you were to build the client as
-    # `client = aws_encryption_sdk.EncryptionSDKClient()`.
-    client = aws_encryption_sdk.EncryptionSDKClient(
-        commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
-    )
-
-    # 2. Create encryption context.
-    # Remember that your encryption context is NOT SECRET.
-    # For more information, see
-    # https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
-    encryption_context: Dict[str, str] = {
-        "encryption": "context",
-        "is not": "secret",
-        "but adds": "useful metadata",
-        "that can help you": "be confident that",
-        "the data you are handling": "is what you think it is",
-    }
-
-    # 3. Create a Raw RSA keyring
-    raw_rsa_keyring = generate_rsa_keyring()
-
-    # 4. Encrypt the data for the encryptionContext
+    # 6. Encrypt the data for the encryptionContext
     ciphertext, _ = client.encrypt(
         source=EXAMPLE_DATA,
         keyring=raw_rsa_keyring,
         encryption_context=encryption_context
     )
 
-    # 5. Demonstrate that the ciphertext and plaintext are different.
+    # 7. Demonstrate that the ciphertext and plaintext are different.
     # (This is an example for demonstration; you do not need to do this in your own code.)
     assert ciphertext != EXAMPLE_DATA, \
         "Ciphertext and plaintext data are the same. Invalid encryption"
 
-    # 6. Decrypt your encrypted data using the same keyring you used on encrypt.
+    # 8. Decrypt your encrypted data using the same keyring you used on encrypt.
     plaintext_bytes, dec_header = client.decrypt(
         source=ciphertext,
         keyring=raw_rsa_keyring
     )
 
-    # 7. Demonstrate that the encryption context is correct in the decrypted message header
+    # 9. Demonstrate that the encryption context is correct in the decrypted message header
     # (This is an example for demonstration; you do not need to do this in your own code.)
     for k, v in encryption_context.items():
         assert v == dec_header.encryption_context[k], \
             "Encryption context does not match expected values"
 
-    # 8. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 10. Demonstrate that the decrypted plaintext is identical to the original plaintext.
     # (This is an example for demonstration; you do not need to do this in your own code.)
     assert plaintext_bytes == EXAMPLE_DATA