chore: updated and refactored raw rsa keyring example

Author: RitvikKapila

examples/src/keyrings/raw_rsa_keyring_example.py

@@ -17,18 +17,18 @@
 1. Ciphertext and plaintext data are not the same
 2. Encryption context is correct in the decrypted message header
 3. Decrypted plaintext value matches EXAMPLE_DATA
-4. After verifying that the encrypt and decrypt works, this example also demonstrates
-that the original ciphertext should not be decrypted using a new Raw RSA keyring generated by
-another user, let's say Bob (Points 9 and 10).
+4. The original ciphertext is not decryptable using a keyring with a different RSA key pair
 These sanity checks are for demonstration in the example only. You do not need these in your code.
 
 A Raw RSA keyring that encrypts and decrypts must include an asymmetric public key and private
 key pair. However, you can encrypt data with a Raw RSA keyring that has only a public key,
-and you can decrypt data with a Raw RSA keyring that has only a private key. You can include
-any Raw RSA keyring in a multi-keyring. If you configure a Raw RSA keyring with a public and
-private key, be sure that they are part of the same key pair. Some language implementations
-of the AWS Encryption SDK will not construct a Raw RSA keyring with keys from different pairs.
-Others rely on you to verify that your keys are from the same key pair.
+and you can decrypt data with a Raw RSA keyring that has only a private key. This example requires
+the user to either provide both private and public keys, or not provide any keys and the example
+generates both to test encryption and decryption. If you configure a Raw RSA keyring with a
+public and private key, be sure that they are part of the same key pair. Some language
+implementations of the AWS Encryption SDK will not construct a Raw RSA keyring with keys
+from different pairs. Others rely on you to verify that your keys are from the same key pair.
+You can include any Raw RSA keyring in a multi-keyring.
 
 For more information on how to use Raw RSA keyrings, see
 https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/use-raw-rsa-keyring.html
@@ -58,20 +58,31 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
-def generate_rsa_keyring():
-    """Generates new public and private keys to create a Raw RSA keyring and
-    then generates the keyring
+def should_generate_new_rsa_key_pair(public_key, private_key):
+    """Returns True if user doesn't provide keys, and we need to generate them and
+    returns False if the user has already provided both public and private keys
+    Raises an AssertionError if the user only provides one of private_key and public_key
 
-    Usage: generate_rsa_keyring()
+    Usage: should_generate_new_rsa_key_pair(public_key, private_key)
     """
-    # 1. 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"
+    # If only one of public_key and private_key is provided, raise an Assertion Error
+    if (public_key and not private_key) or (not public_key and private_key):
+        raise AssertionError("Either both public and private keys should be provided! Or no keys \
+                             should be provided and the example can create the keys for you!")
+
+    # If no keys are provided, we should generate a new rsa key pair, so return True
+    if not public_key and not private_key:
+        return True
+
+    # If both keys are already provided, return False
+    return False
+
 
-    # 2. Generate a 4096-bit RSA key to use with your keyring.
+def generate_rsa_keys():
+    """Generates a 4096-bit RSA public and private key pair
+
+    Usage: generate_rsa_keys()
+    """
     ssh_rsa_exponent = 65537
     bit_strength = 4096
     key = rsa.generate_private_key(
@@ -81,19 +92,34 @@ def generate_rsa_keyring():
     )
 
     # This example choses a particular type of encoding, format and encryption_algorithm
-    # Users can choose the PrivateFormat, PublicFormat and encryption_algorithm that align most
+    # Users can choose the PublicFormat, PrivateFormat and encryption_algorithm that align most
     # with their use-cases
+    public_key = key.public_key().public_bytes(
+        encoding=crypto_serialization.Encoding.PEM,
+        format=crypto_serialization.PublicFormat.SubjectPublicKeyInfo
+    )
     private_key = key.private_bytes(
         encoding=crypto_serialization.Encoding.PEM,
         format=crypto_serialization.PrivateFormat.TraditionalOpenSSL,
         encryption_algorithm=crypto_serialization.NoEncryption()
     )
-    public_key = key.public_key().public_bytes(
-        encoding=crypto_serialization.Encoding.PEM,
-        format=crypto_serialization.PublicFormat.SubjectPublicKeyInfo
-    )
 
-    # 3. Create a Raw RSA keyring
+    return public_key, private_key
+
+
+def create_rsa_keyring(public_key, private_key):
+    """Create a Raw RSA keyring using the key pair
+
+    Usage: create_rsa_keyring(public_key, private_key)
+    """
+    # 1. 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. Create a Raw RSA keyring
     mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
         config=MaterialProvidersConfig()
     )
@@ -113,10 +139,12 @@ def generate_rsa_keyring():
     return raw_rsa_keyring
 
 
-def encrypt_and_decrypt_with_keyring():
-    """Demonstrate an encrypt/decrypt cycle using a Raw RSA keyring.
+def encrypt_and_decrypt_with_keyring(public_key=None, private_key=None):
+    """Demonstrate an encrypt/decrypt cycle using a Raw RSA keyring
+    with user defined keys. If no keys are present, generate new RSA
+    public and private keys and use them to create a Raw RSA keyring
 
-    Usage: encrypt_and_decrypt_with_keyring()
+    Usage: encrypt_and_decrypt_with_keyring(public_key, private_key)
     """
     # 1. Instantiate the encryption SDK client.
     # This builds the client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy,
@@ -142,7 +170,18 @@ def encrypt_and_decrypt_with_keyring():
     }
 
     # 3. Create a Raw RSA keyring
-    raw_rsa_keyring = generate_rsa_keyring()
+
+    # Check if we need to generate an RSA key pair
+    should_generate_new_rsa_key_pair_bool = \
+        should_generate_new_rsa_key_pair(public_key=public_key, private_key=private_key)
+
+    # If user doesn't provide the keys, that is, if should_generate_new_rsa_key_pair_bool is True
+    # generate a new RSA public and private key pair
+    if should_generate_new_rsa_key_pair_bool:
+        public_key, private_key = generate_rsa_keys()
+
+    # Create the keyring
+    raw_rsa_keyring = create_rsa_keyring(public_key=public_key, private_key=private_key)
 
     # 4. Encrypt the data for the encryptionContext
     ciphertext, _ = client.encrypt(
@@ -176,12 +215,16 @@ def encrypt_and_decrypt_with_keyring():
     # decryption of the original ciphertext is not possible with a different keyring (Bob's)
     # (This is an example for demonstration; you do not need to do this in your own code.)
 
-    # 9. Generate a new Raw RSA keyring for Bob
-    raw_rsa_keyring_bob = generate_rsa_keyring()
+    # 9. Create a new Raw RSA keyring for Bob
+    # Generate new keys
+    public_key_bob, private_key_bob = generate_rsa_keys()
+
+    # Create the keyring
+    raw_rsa_keyring_bob = create_rsa_keyring(public_key=public_key_bob, private_key=private_key_bob)
 
     # 10. Test decrypt for the original ciphertext using raw_rsa_keyring_bob
     try:
-        plaintext_bytes_bob, dec_header_bob = client.decrypt(
+        plaintext_bytes_bob, _ = client.decrypt(
             source=ciphertext,
             keyring=raw_rsa_keyring_bob
         )

examples/test/keyrings/test_i_raw_rsa_keyring_example.py

@@ -3,11 +3,48 @@
 """Test suite for the Raw RSA keyring example."""
 import pytest
 
-from ...src.keyrings.raw_rsa_keyring_example import encrypt_and_decrypt_with_keyring
+from ...src.keyrings.raw_rsa_keyring_example import encrypt_and_decrypt_with_keyring, generate_rsa_keys
 
 pytestmark = [pytest.mark.examples]
 
 
-def test_encrypt_and_decrypt_with_keyring():
-    """Test function for encrypt and decrypt using the Raw RSA Keyring example."""
+def test_encrypt_and_decrypt_with_keyring_without_user_defined_keys():
+    """Test function for encrypt and decrypt using the Raw RSA Keyring example
+    where user doesn't provide the public and private keys"""
     encrypt_and_decrypt_with_keyring()
+
+
+def test_encrypt_and_decrypt_with_keyring_with_user_defined_keys():
+    """Test function for encrypt and decrypt using the Raw RSA Keyring example
+    where user provides the public and private keys. To test this, we create the
+    keys using the generate_rsa_keys function"""
+    user_public_key, user_private_key = generate_rsa_keys()
+    encrypt_and_decrypt_with_keyring(public_key=user_public_key, private_key=user_private_key)
+
+
+def test_encrypt_and_decrypt_fails_if_user_provides_only_public_key():
+    """Test function for encrypt and decrypt using the Raw RSA Keyring example
+    where user provides only the public key. The program should throw an Assertion error
+    as this example requires the user to either provide both private and public keys to
+    test both encryption and decryption, or not provide any keys and the example generates both"""
+    user_public_key, user_private_key = generate_rsa_keys()
+    try:
+        encrypt_and_decrypt_with_keyring(public_key=user_public_key)
+
+        raise AssertionError("encrypt_and_decrypt_with_keyring should raise an error")
+    except AssertionError:
+        pass
+
+
+def test_encrypt_and_decrypt_fails_if_user_provides_only_private_key():
+    """Test function for encrypt and decrypt using the Raw RSA Keyring example
+    where user provides only the private key. The program should throw an Assertion error
+    as this example requires the user to either provide both private and public keys to
+    test both encryption and decryption, or not provide any keys and the example generates both"""
+    user_public_key, user_private_key = generate_rsa_keys()
+    try:
+        encrypt_and_decrypt_with_keyring(private_key=user_private_key)
+
+        raise AssertionError("encrypt_and_decrypt_with_keyring should raise an error")
+    except AssertionError:
+        pass