refactoring

Author: RitvikKapila

examples/src/multithreading/__init__.py

@@ -0,0 +1,45 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""init file for multi-threading examples."""
+import aws_encryption_sdk
+from aws_cryptographic_materialproviders.mpl.references import IKeyring
+from typing import Dict  # noqa pylint: disable=wrong-import-order
+
+
+def encrypt_and_decrypt_with_keyring(
+    plaintext_data: bytes,
+    keyring: IKeyring,
+    esdk_client: aws_encryption_sdk.EncryptionSDKClient
+):
+    """encrypt_and_decrypt_with_keyring how to encrypt plaintext data using a Raw AES keyring.
+
+    Usage: encrypt_and_decrypt_with_keyring(plaintext_data, keyring, esdk_client)
+    :param plaintext_data: plaintext data you want to encrypt
+    :type: bytes
+    :param keyring: Keyring to use for encryption.
+    :type keyring: IKeyring
+    :param esdk_client: The Encryption SDK client to use for encryption.
+    :type esdk_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, _ = esdk_client.encrypt(
+        source=plaintext_data,
+        keyring=keyring,
+        encryption_context=encryption_context
+    )
+
+    decrypted_plaintext_data, _ = esdk_client.decrypt(
+        source=ciphertext_data,
+        keyring=keyring
+    )
+
+    return decrypted_plaintext_data

examples/src/multithreading/raw_aes_keyring.py

@@ -0,0 +1,37 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""This file contains methdos to use for the multi-threaded Raw AES keyring."""
+
+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
+import secrets
+
+
+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,73 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""This file contains methdos to use for the multi-threaded 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 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
+from typing import Dict  # noqa pylint: disable=wrong-import-order
+
+
+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/multithreaded/test_i_raw_aes_keyring_multithreaded_example.py

@@ -0,0 +1,33 @@
+# 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
+
+from ...src.multithreading.raw_aes_keyring import create_keyring
+from ...src.multithreading import encrypt_and_decrypt_with_keyring
+
+import aws_encryption_sdk
+from aws_encryption_sdk import CommitmentPolicy
+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 AES Keyring example."""
+    keyring = create_keyring()
+    plaintext_data = b"Hello World"
+    esdk_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,
+                                          esdk_client=esdk_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"

examples/test/multithreaded/test_i_raw_rsa_keyring_multithreaded_example.py

@@ -0,0 +1,33 @@
+# 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
+from ...src.multithreading.raw_rsa_keyring import generate_rsa_keys, create_keyring
+from ...src.multithreading import encrypt_and_decrypt_with_keyring
+import aws_encryption_sdk
+from aws_encryption_sdk import CommitmentPolicy
+
+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"
+    esdk_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,
+                                          esdk_client=esdk_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"