diff --git a/examples/src/multithreading/__init__.py b/examples/src/multithreading/__init__.py new file mode 100644 index 000000000..32210a0ab --- /dev/null +++ b/examples/src/multithreading/__init__.py @@ -0,0 +1,63 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""init file for multi-threading examples.""" +import time + +from aws_cryptographic_materialproviders.mpl.references import IKeyring +from typing import Dict # noqa pylint: disable=wrong-import-order + +import aws_encryption_sdk + + +def encrypt_and_decrypt_with_keyring( + plaintext_data: bytes, + keyring: IKeyring, + client: aws_encryption_sdk.EncryptionSDKClient +): + """Demonstrate how to encrypt and decrypt plaintext data using a keyring. + + Usage: encrypt_and_decrypt_with_keyring(plaintext_data, keyring, client) + :param plaintext_data: plaintext data you want to encrypt + :type: bytes + :param keyring: Keyring to use for encryption. + :type keyring: IKeyring + :param client: The Encryption SDK client to use for encryption. + :type client: aws_encryption_sdk.EncryptionSDKClient + :return: encrypted and decrypted (cycled) plaintext data + :rtype: bytes + """ + 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", + } + + ciphertext_data, _ = client.encrypt( + source=plaintext_data, + keyring=keyring, + encryption_context=encryption_context + ) + + decrypted_plaintext_data, _ = client.decrypt( + source=ciphertext_data, + keyring=keyring + ) + + return decrypted_plaintext_data + + +def run_encrypt_and_decrypt_with_keyring_for_duration_seconds( + plaintext_data: bytes, + keyring: IKeyring, + client: aws_encryption_sdk.EncryptionSDKClient, + duration: int = 2 +): + """Helper function to repeatedly run an encrypt and decrypt cycle for 'duration' seconds.""" + time_end = time.time() + duration + + while time.time() < time_end: + decrypted_plaintext_data = encrypt_and_decrypt_with_keyring(plaintext_data, keyring, client) + assert decrypted_plaintext_data == plaintext_data, \ + "Decrypted plaintext should be identical to the original plaintext. Invalid decryption" diff --git a/examples/src/multithreading/raw_aes_keyring.py b/examples/src/multithreading/raw_aes_keyring.py new file mode 100644 index 000000000..42a391344 --- /dev/null +++ b/examples/src/multithreading/raw_aes_keyring.py @@ -0,0 +1,38 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This file contains methods to use for testing multi-threading for Raw AES keyring.""" + +import secrets + +from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders +from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig +from aws_cryptographic_materialproviders.mpl.models import AesWrappingAlg, CreateRawAesKeyringInput +from aws_cryptographic_materialproviders.mpl.references import IKeyring + + +def create_keyring(): + """Demonstrate how to create a Raw AES keyring. + + Usage: create_keyring() + """ + key_name_space = "Some managed raw keys" + key_name = "My 256-bit AES wrapping key" + + static_key = secrets.token_bytes(32) + + mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders( + config=MaterialProvidersConfig() + ) + + keyring_input: CreateRawAesKeyringInput = CreateRawAesKeyringInput( + key_namespace=key_name_space, + key_name=key_name, + wrapping_key=static_key, + wrapping_alg=AesWrappingAlg.ALG_AES256_GCM_IV12_TAG16 + ) + + keyring: IKeyring = mat_prov.create_raw_aes_keyring( + input=keyring_input + ) + + return keyring diff --git a/examples/src/multithreading/raw_rsa_keyring.py b/examples/src/multithreading/raw_rsa_keyring.py new file mode 100644 index 000000000..77d26f2fc --- /dev/null +++ b/examples/src/multithreading/raw_rsa_keyring.py @@ -0,0 +1,66 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This file contains methods to use for testing multi-threading for Raw RSA keyring.""" +from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders +from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig +from aws_cryptographic_materialproviders.mpl.models import CreateRawRsaKeyringInput, PaddingScheme +from aws_cryptographic_materialproviders.mpl.references import IKeyring +from cryptography.hazmat.backends import default_backend as crypto_default_backend +from cryptography.hazmat.primitives import serialization as crypto_serialization +from cryptography.hazmat.primitives.asymmetric import rsa + + +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( + backend=crypto_default_backend(), + public_exponent=ssh_rsa_exponent, + key_size=bit_strength + ) + + # This example choses a particular type of encoding, format and encryption_algorithm + # 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() + ) + + return public_key, private_key + + +def create_keyring(public_key, private_key): + """Demonstrate how to create a Raw RSA keyring using the key pair. + + Usage: create_keyring(public_key, private_key) + """ + key_name_space = "Some managed raw keys" + key_name = "My 4096-bit RSA wrapping key" + + mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders( + config=MaterialProvidersConfig() + ) + + keyring_input: CreateRawRsaKeyringInput = CreateRawRsaKeyringInput( + key_namespace=key_name_space, + key_name=key_name, + padding_scheme=PaddingScheme.OAEP_SHA256_MGF1, + public_key=public_key, + private_key=private_key + ) + + keyring: IKeyring = mat_prov.create_raw_rsa_keyring( + input=keyring_input + ) + + return keyring diff --git a/examples/test/multithreading/__init__.py b/examples/test/multithreading/__init__.py new file mode 100644 index 000000000..120179eda --- /dev/null +++ b/examples/test/multithreading/__init__.py @@ -0,0 +1,3 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Stub module indicator to make linter configuration simpler.""" diff --git a/examples/test/multithreading/test_i_raw_aes_keyring_multithreaded_example.py b/examples/test/multithreading/test_i_raw_aes_keyring_multithreaded_example.py new file mode 100644 index 000000000..225e3c411 --- /dev/null +++ b/examples/test/multithreading/test_i_raw_aes_keyring_multithreaded_example.py @@ -0,0 +1,50 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Test suite for the Raw AES keyring example with multi-threading.""" +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pytest +# pylint and isort disagree about where this goes; listen to isort +from typing import Optional # pylint: disable=wrong-import-order + +import aws_encryption_sdk +from aws_encryption_sdk import CommitmentPolicy + +from ...src.multithreading import run_encrypt_and_decrypt_with_keyring_for_duration_seconds +from ...src.multithreading.raw_aes_keyring import create_keyring + +pytestmark = [pytest.mark.examples] + + +def encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=64, duration=60): + """Encrypt and decrypt using a keyring for fixed n_threads and duration.""" + keyring = create_keyring() + plaintext_data = b"Hello World" + client = aws_encryption_sdk.EncryptionSDKClient( + commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT + ) + + with ThreadPoolExecutor(max_workers=n_threads) as executor: + thread_futures = {executor.submit(run_encrypt_and_decrypt_with_keyring_for_duration_seconds, + plaintext_data=plaintext_data, + keyring=keyring, + client=client, + duration=duration): i for i in range(n_threads)} + + for future in as_completed(thread_futures): + future.result() + + +def test_encrypt_and_decrypt_with_keyring_multithreaded( + n_threads_list: Optional[list] = None, + duration_list: Optional[list] = None, +): + """Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration.""" + # Set defaults if no value is provided + if n_threads_list is None: + n_threads_list = [1, 4, 16, 64] + if duration_list is None: + duration_list = [2, 10, 60] + for n in n_threads_list: + for d in duration_list: + encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=n, duration=d) diff --git a/examples/test/multithreading/test_i_raw_rsa_keyring_multithreaded_example.py b/examples/test/multithreading/test_i_raw_rsa_keyring_multithreaded_example.py new file mode 100644 index 000000000..c1ebdbe24 --- /dev/null +++ b/examples/test/multithreading/test_i_raw_rsa_keyring_multithreaded_example.py @@ -0,0 +1,51 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Test suite for the Raw RSA keyring example with multi-threading.""" +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pytest +# pylint and isort disagree about where this goes; listen to isort +from typing import Optional # pylint: disable=wrong-import-order + +import aws_encryption_sdk +from aws_encryption_sdk import CommitmentPolicy + +from ...src.multithreading import run_encrypt_and_decrypt_with_keyring_for_duration_seconds +from ...src.multithreading.raw_rsa_keyring import create_keyring, generate_rsa_keys + +pytestmark = [pytest.mark.examples] + + +def encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=64, duration=60): + """Encrypt and decrypt using a keyring for fixed n_threads and duration.""" + public_key, private_key = generate_rsa_keys() + keyring = create_keyring(public_key=public_key, private_key=private_key) + plaintext_data = b"Hello World" + client = aws_encryption_sdk.EncryptionSDKClient( + commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT + ) + + with ThreadPoolExecutor(max_workers=n_threads) as executor: + thread_futures = {executor.submit(run_encrypt_and_decrypt_with_keyring_for_duration_seconds, + plaintext_data=plaintext_data, + keyring=keyring, + client=client, + duration=duration): i for i in range(n_threads)} + + for future in as_completed(thread_futures): + future.result() + + +def test_encrypt_and_decrypt_with_keyring_multithreaded( + n_threads_list: Optional[list] = None, + duration_list: Optional[list] = None, +): + """Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration.""" + # Set defaults if no value is provided + if n_threads_list is None: + n_threads_list = [1, 4, 16, 64] + if duration_list is None: + duration_list = [2, 10, 60] + for n in n_threads_list: + for d in duration_list: + encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=n, duration=d)