add

Author: lucasmcdonald3

examples/src/multithreading/__init__.py

@@ -0,0 +1,62 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""init file for multi-threading examples."""
+from aws_cryptographic_materialproviders.mpl.references import IKeyring
+from typing import Dict  # noqa pylint: disable=wrong-import-order
+
+import aws_encryption_sdk
+import time
+
+
+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"

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

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

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."""

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
+
+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
+
+import time
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=16, duration=60):
+    """Helper function for multi-threaded encrypt and decrypt using a keyring for fixed n_threads and duration."""
+    print(n_threads, duration)
+    start_time = time.time()
+    print('start_time', start_time)
+    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()
+    end_time = time.time()
+    print('end_time', end_time)
+    print('duration', end_time - start_time)
+
+
+# def test_encrypt_and_decrypt_with_keyring_multithreaded(n_threads_list: list = [4, 16, 64], duration_list: list = [2, 10, 60]):
+#     """Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration."""
+#     print('hello', n_threads_list, duration_list)
+#     for n in n_threads_list:
+#         for d in duration_list:
+#             print(n, d, time.time())
+#             encrypt_and_decrypt_with_keyring_helper(n_threads=n, duration=d)

examples/test/multithreading/test_i_raw_rsa_keyring_multithreaded_example.py

@@ -0,0 +1,35 @@
+# 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
+
+import aws_encryption_sdk
+from aws_encryption_sdk import CommitmentPolicy
+
+from ...src.multithreading import encrypt_and_decrypt_with_keyring
+from ...src.multithreading.raw_rsa_keyring import create_keyring, generate_rsa_keys
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_encrypt_and_decrypt_with_keyring(n_threads=10):
+    """Test function for multi-threaded encrypt and decrypt using the Raw RSA Keyring example."""
+    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(encrypt_and_decrypt_with_keyring,
+                                          plaintext_data=plaintext_data,
+                                          keyring=keyring,
+                                          client=client): i for i in range(n_threads)}
+
+        for future in as_completed(thread_futures):
+            decrypted_plaintext_data = future.result()
+            assert decrypted_plaintext_data == plaintext_data, \
+                "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"