updated raw rsa keyring to get keys from user files

Author: RitvikKapila

.gitignore

@@ -31,6 +31,7 @@ __pycache__
 
 # PyTest
 .pytest_cache
+test_keys/
 
 # PyCharm
 .idea/

examples/src/keyrings/raw_aes_keyring_example.py

@@ -79,6 +79,8 @@ def encrypt_and_decrypt_with_keyring():
     }
 
     # 4. Generate a 256-bit AES key to use with your keyring.
+    # In practice, you should get this key from a secure key management system such as an HSM.
+
     # Here, the input to secrets.token_bytes() = 32 bytes = 256 bits
     static_key = secrets.token_bytes(32)
 
@@ -98,7 +100,7 @@ def encrypt_and_decrypt_with_keyring():
         input=keyring_input
     )
 
-    # 6. Encrypt the data for the encryptionContext
+    # 6. Encrypt the data with the encryptionContext
     ciphertext, _ = client.encrypt(
         source=EXAMPLE_DATA,
         keyring=raw_aes_keyring,

examples/src/keyrings/raw_rsa_keyring_example.py

@@ -58,20 +58,21 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
-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
+def should_generate_new_rsa_key_pair(public_key_file_name, private_key_file_name):
+    """Returns True if user doesn't provide keys, and we need to generate them;
+    Returns False if the user has already provided both public and private keys
+    Raises a ValueError if the user only provides one of private_key and public_key
 
-    Usage: should_generate_new_rsa_key_pair(public_key, private_key)
+    Usage: should_generate_new_rsa_key_pair(public_key_file_name, private_key_file_name)
     """
-    # 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 \
+    # If only one of public_key and private_key files is provided, raise a ValueError
+    if (public_key_file_name and not private_key_file_name)\
+            or (not public_key_file_name and private_key_file_name):
+        raise ValueError("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:
+    if not public_key_file_name and not private_key_file_name:
         return True
 
     # If both keys are already provided, return False
@@ -139,12 +140,12 @@ def create_rsa_keyring(public_key, private_key):
     return raw_rsa_keyring
 
 
-def encrypt_and_decrypt_with_keyring(public_key=None, private_key=None):
+def encrypt_and_decrypt_with_keyring(public_key_file_name=None, private_key_file_name=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(public_key, private_key)
+    Usage: encrypt_and_decrypt_with_keyring(public_key_file_name, private_key_file_name)
     """
     # 1. Instantiate the encryption SDK client.
     # This builds the client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy,
@@ -173,17 +174,31 @@ def encrypt_and_decrypt_with_keyring(public_key=None, private_key=None):
 
     # 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)
+        should_generate_new_rsa_key_pair(public_key_file_name=public_key_file_name,
+                                         private_key_file_name=private_key_file_name)
 
     # 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()
+    else:
+        # If user provides the keys, read the keys from the files
+        with open(public_key_file_name, "r", encoding='utf-8') as f:
+            public_key = f.read()
+
+        # Convert the public key from a string to bytes
+        public_key = bytes(public_key, 'utf-8')
+
+        with open(private_key_file_name, "r", encoding='utf-8') as f:
+            private_key = f.read()
+
+        # Convert the private key from a string to bytes
+        private_key = bytes(private_key, 'utf-8')
 
     # Create the keyring
     raw_rsa_keyring = create_rsa_keyring(public_key=public_key, private_key=private_key)
 
-    # 4. Encrypt the data for the encryptionContext
+    # 4. Encrypt the data with the encryptionContext
     ciphertext, _ = client.encrypt(
         source=EXAMPLE_DATA,
         keyring=raw_rsa_keyring,
@@ -212,7 +227,7 @@ def encrypt_and_decrypt_with_keyring(public_key=None, private_key=None):
     assert plaintext_bytes == EXAMPLE_DATA
 
     # The next part of the example creates a new RSA keyring (for Bob) to demonstrate that
-    # decryption of the original ciphertext is not possible with a different keyring (Bob's)
+    # 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. Create a new Raw RSA keyring for Bob

examples/test/keyrings/test_i_raw_rsa_keyring_example.py

@@ -1,6 +1,8 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Test suite for the Raw RSA keyring example."""
+import os
+
 import pytest
 
 from ...src.keyrings.raw_rsa_keyring_example import encrypt_and_decrypt_with_keyring, generate_rsa_keys
@@ -20,39 +22,95 @@ def test_encrypt_and_decrypt_with_keyring_with_user_defined_keys():
     """Test function for encrypt and decrypt using the Raw RSA Keyring example.
 
     Here user provides the public and private keys. To test this, we create the
-    keys using the generate_rsa_keys function
+    keys using the generate_rsa_keys function and write them to the file.
+    Then we call the encrypt_and_decrypt_with_keyring function and pass them
     """
+    # Generate the user keys for testing
     user_public_key, user_private_key = generate_rsa_keys()
-    encrypt_and_decrypt_with_keyring(public_key=user_public_key, private_key=user_private_key)
+
+    # Convert the keys to strings
+    user_public_key = user_public_key.decode('utf-8')
+    user_private_key = user_private_key.decode('utf-8')
+
+    test_keys_directory = 'test_keys'
+    if not os.path.exists(test_keys_directory):
+        os.makedirs(test_keys_directory)
+
+    # Define the file names for the keys
+    user_public_key_file_name = test_keys_directory + '/user_public_key_file_name.pem'
+    user_private_key_file_name = test_keys_directory + '/user_private_key_file_name.pem'
+
+    # Write the public key to the file
+    with open(user_public_key_file_name, "w", encoding="utf-8") as f:
+        f.write(user_public_key)
+
+    # Write the private key to the file
+    with open(user_private_key_file_name, "w", encoding="utf-8") as f:
+        f.write(user_private_key)
+
+    encrypt_and_decrypt_with_keyring(public_key_file_name=user_public_key_file_name,
+                                     private_key_file_name=user_private_key_file_name)
 
 
 def test_encrypt_and_decrypt_fails_if_user_provides_only_public_key():
     """Test function for encrypt and decrypt using the Raw RSA Keyring example.
 
-    Here user provides only the public key. The program should throw an Assertion error
+    Here user provides only the public key. The program should throw an Value 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
     """
+    # Generate the user keys for testing
     user_public_key, user_private_key = generate_rsa_keys()
+
+    # Convert the public key to string
+    user_public_key = user_public_key.decode('utf-8')
+
+    test_keys_directory = 'test_keys'
+    if not os.path.exists(test_keys_directory):
+        os.makedirs(test_keys_directory)
+
+    # Define the file name for the public key
+    user_public_key_file_name = test_keys_directory + '/user_public_key_file_name.pem'
+
+    # Write the public key to the file
+    with open(user_public_key_file_name, "w", encoding="utf-8") as f:
+        f.write(user_public_key)
+
     try:
-        encrypt_and_decrypt_with_keyring(public_key=user_public_key)
+        encrypt_and_decrypt_with_keyring(public_key_file_name=user_public_key_file_name)
 
         raise AssertionError("encrypt_and_decrypt_with_keyring should raise an error")
-    except AssertionError:
+    except ValueError:
         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.
 
-    Here user provides only the private key. The program should throw an Assertion error
+    Here user provides only the private key. The program should throw an Value 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
     """
+    # Generate the user keys for testing
     user_public_key, user_private_key = generate_rsa_keys()
+
+    # Convert the private key to string
+    user_private_key = user_private_key.decode('utf-8')
+
+    test_keys_directory = 'test_keys'
+    if not os.path.exists(test_keys_directory):
+        os.makedirs(test_keys_directory)
+
+    # Define the file name for the private key
+    user_private_key_file_name = test_keys_directory + '/user_private_key_file_name.pem'
+
+    # Write the private key to the file
+    with open(user_private_key_file_name, "w", encoding="utf-8") as f:
+        f.write(user_private_key)
+
     try:
-        encrypt_and_decrypt_with_keyring(private_key=user_private_key)
+        encrypt_and_decrypt_with_keyring(private_key_file_name=user_private_key_file_name)
 
         raise AssertionError("encrypt_and_decrypt_with_keyring should raise an error")
-    except AssertionError:
+    except ValueError:
         pass