Skip to content

Commit 5e813e7

Browse files
lucasmcdonald3lucasmcdonald3
authored
chore(examples): Added raw RSA/AES keyring multithreaded examples (#694)
1 parent 0de58cd commit 5e813e7

File tree

6 files changed

+271
-0
lines changed

6 files changed

+271
-0
lines changed

‎examples/src/multithreading/__init__.py

+63
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""init file for multi-threading examples."""
4+
import time
5+
6+
from aws_cryptographic_materialproviders.mpl.references import IKeyring
7+
from typing import Dict # noqa pylint: disable=wrong-import-order
8+
9+
import aws_encryption_sdk
10+
11+
12+
def encrypt_and_decrypt_with_keyring(
13+
plaintext_data: bytes,
14+
keyring: IKeyring,
15+
client: aws_encryption_sdk.EncryptionSDKClient
16+
):
17+
"""Demonstrate how to encrypt and decrypt plaintext data using a keyring.
18+
19+
Usage: encrypt_and_decrypt_with_keyring(plaintext_data, keyring, client)
20+
:param plaintext_data: plaintext data you want to encrypt
21+
:type: bytes
22+
:param keyring: Keyring to use for encryption.
23+
:type keyring: IKeyring
24+
:param client: The Encryption SDK client to use for encryption.
25+
:type client: aws_encryption_sdk.EncryptionSDKClient
26+
:return: encrypted and decrypted (cycled) plaintext data
27+
:rtype: bytes
28+
"""
29+
encryption_context: Dict[str, str] = {
30+
"encryption": "context",
31+
"is not": "secret",
32+
"but adds": "useful metadata",
33+
"that can help you": "be confident that",
34+
"the data you are handling": "is what you think it is",
35+
}
36+
37+
ciphertext_data, _ = client.encrypt(
38+
source=plaintext_data,
39+
keyring=keyring,
40+
encryption_context=encryption_context
41+
)
42+
43+
decrypted_plaintext_data, _ = client.decrypt(
44+
source=ciphertext_data,
45+
keyring=keyring
46+
)
47+
48+
return decrypted_plaintext_data
49+
50+
51+
def run_encrypt_and_decrypt_with_keyring_for_duration_seconds(
52+
plaintext_data: bytes,
53+
keyring: IKeyring,
54+
client: aws_encryption_sdk.EncryptionSDKClient,
55+
duration: int = 2
56+
):
57+
"""Helper function to repeatedly run an encrypt and decrypt cycle for 'duration' seconds."""
58+
time_end = time.time() + duration
59+
60+
while time.time() < time_end:
61+
decrypted_plaintext_data = encrypt_and_decrypt_with_keyring(plaintext_data, keyring, client)
62+
assert decrypted_plaintext_data == plaintext_data, \
63+
"Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

‎examples/src/multithreading/raw_aes_keyring.py

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""This file contains methods to use for testing multi-threading for Raw AES keyring."""
4+
5+
import secrets
6+
7+
from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
8+
from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
9+
from aws_cryptographic_materialproviders.mpl.models import AesWrappingAlg, CreateRawAesKeyringInput
10+
from aws_cryptographic_materialproviders.mpl.references import IKeyring
11+
12+
13+
def create_keyring():
14+
"""Demonstrate how to create a Raw AES keyring.
15+
16+
Usage: create_keyring()
17+
"""
18+
key_name_space = "Some managed raw keys"
19+
key_name = "My 256-bit AES wrapping key"
20+
21+
static_key = secrets.token_bytes(32)
22+
23+
mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
24+
config=MaterialProvidersConfig()
25+
)
26+
27+
keyring_input: CreateRawAesKeyringInput = CreateRawAesKeyringInput(
28+
key_namespace=key_name_space,
29+
key_name=key_name,
30+
wrapping_key=static_key,
31+
wrapping_alg=AesWrappingAlg.ALG_AES256_GCM_IV12_TAG16
32+
)
33+
34+
keyring: IKeyring = mat_prov.create_raw_aes_keyring(
35+
input=keyring_input
36+
)
37+
38+
return keyring

‎examples/src/multithreading/raw_rsa_keyring.py

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""This file contains methods to use for testing multi-threading for Raw RSA keyring."""
4+
from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
5+
from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
6+
from aws_cryptographic_materialproviders.mpl.models import CreateRawRsaKeyringInput, PaddingScheme
7+
from aws_cryptographic_materialproviders.mpl.references import IKeyring
8+
from cryptography.hazmat.backends import default_backend as crypto_default_backend
9+
from cryptography.hazmat.primitives import serialization as crypto_serialization
10+
from cryptography.hazmat.primitives.asymmetric import rsa
11+
12+
13+
def generate_rsa_keys():
14+
"""Generates a 4096-bit RSA public and private key pair
15+
16+
Usage: generate_rsa_keys()
17+
"""
18+
ssh_rsa_exponent = 65537
19+
bit_strength = 4096
20+
key = rsa.generate_private_key(
21+
backend=crypto_default_backend(),
22+
public_exponent=ssh_rsa_exponent,
23+
key_size=bit_strength
24+
)
25+
26+
# This example choses a particular type of encoding, format and encryption_algorithm
27+
# Users can choose the PublicFormat, PrivateFormat and encryption_algorithm that align most
28+
# with their use-cases
29+
public_key = key.public_key().public_bytes(
30+
encoding=crypto_serialization.Encoding.PEM,
31+
format=crypto_serialization.PublicFormat.SubjectPublicKeyInfo
32+
)
33+
private_key = key.private_bytes(
34+
encoding=crypto_serialization.Encoding.PEM,
35+
format=crypto_serialization.PrivateFormat.TraditionalOpenSSL,
36+
encryption_algorithm=crypto_serialization.NoEncryption()
37+
)
38+
39+
return public_key, private_key
40+
41+
42+
def create_keyring(public_key, private_key):
43+
"""Demonstrate how to create a Raw RSA keyring using the key pair.
44+
45+
Usage: create_keyring(public_key, private_key)
46+
"""
47+
key_name_space = "Some managed raw keys"
48+
key_name = "My 4096-bit RSA wrapping key"
49+
50+
mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
51+
config=MaterialProvidersConfig()
52+
)
53+
54+
keyring_input: CreateRawRsaKeyringInput = CreateRawRsaKeyringInput(
55+
key_namespace=key_name_space,
56+
key_name=key_name,
57+
padding_scheme=PaddingScheme.OAEP_SHA256_MGF1,
58+
public_key=public_key,
59+
private_key=private_key
60+
)
61+
62+
keyring: IKeyring = mat_prov.create_raw_rsa_keyring(
63+
input=keyring_input
64+
)
65+
66+
return keyring

‎examples/test/multithreading/__init__.py

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Stub module indicator to make linter configuration simpler."""

‎examples/test/multithreading/test_i_raw_aes_keyring_multithreaded_example.py

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Test suite for the Raw AES keyring example with multi-threading."""
4+
from concurrent.futures import ThreadPoolExecutor, as_completed
5+
6+
import pytest
7+
# pylint and isort disagree about where this goes; listen to isort
8+
from typing import Optional # pylint: disable=wrong-import-order
9+
10+
import aws_encryption_sdk
11+
from aws_encryption_sdk import CommitmentPolicy
12+
13+
from ...src.multithreading import run_encrypt_and_decrypt_with_keyring_for_duration_seconds
14+
from ...src.multithreading.raw_aes_keyring import create_keyring
15+
16+
pytestmark = [pytest.mark.examples]
17+
18+
19+
def encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=64, duration=60):
20+
"""Encrypt and decrypt using a keyring for fixed n_threads and duration."""
21+
keyring = create_keyring()
22+
plaintext_data = b"Hello World"
23+
client = aws_encryption_sdk.EncryptionSDKClient(
24+
commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
25+
)
26+
27+
with ThreadPoolExecutor(max_workers=n_threads) as executor:
28+
thread_futures = {executor.submit(run_encrypt_and_decrypt_with_keyring_for_duration_seconds,
29+
plaintext_data=plaintext_data,
30+
keyring=keyring,
31+
client=client,
32+
duration=duration): i for i in range(n_threads)}
33+
34+
for future in as_completed(thread_futures):
35+
future.result()
36+
37+
38+
def test_encrypt_and_decrypt_with_keyring_multithreaded(
39+
n_threads_list: Optional[list] = None,
40+
duration_list: Optional[list] = None,
41+
):
42+
"""Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration."""
43+
# Set defaults if no value is provided
44+
if n_threads_list is None:
45+
n_threads_list = [1, 4, 16, 64]
46+
if duration_list is None:
47+
duration_list = [2, 10, 60]
48+
for n in n_threads_list:
49+
for d in duration_list:
50+
encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=n, duration=d)

‎examples/test/multithreading/test_i_raw_rsa_keyring_multithreaded_example.py

+51
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Test suite for the Raw RSA keyring example with multi-threading."""
4+
from concurrent.futures import ThreadPoolExecutor, as_completed
5+
6+
import pytest
7+
# pylint and isort disagree about where this goes; listen to isort
8+
from typing import Optional # pylint: disable=wrong-import-order
9+
10+
import aws_encryption_sdk
11+
from aws_encryption_sdk import CommitmentPolicy
12+
13+
from ...src.multithreading import run_encrypt_and_decrypt_with_keyring_for_duration_seconds
14+
from ...src.multithreading.raw_rsa_keyring import create_keyring, generate_rsa_keys
15+
16+
pytestmark = [pytest.mark.examples]
17+
18+
19+
def encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=64, duration=60):
20+
"""Encrypt and decrypt using a keyring for fixed n_threads and duration."""
21+
public_key, private_key = generate_rsa_keys()
22+
keyring = create_keyring(public_key=public_key, private_key=private_key)
23+
plaintext_data = b"Hello World"
24+
client = aws_encryption_sdk.EncryptionSDKClient(
25+
commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
26+
)
27+
28+
with ThreadPoolExecutor(max_workers=n_threads) as executor:
29+
thread_futures = {executor.submit(run_encrypt_and_decrypt_with_keyring_for_duration_seconds,
30+
plaintext_data=plaintext_data,
31+
keyring=keyring,
32+
client=client,
33+
duration=duration): i for i in range(n_threads)}
34+
35+
for future in as_completed(thread_futures):
36+
future.result()
37+
38+
39+
def test_encrypt_and_decrypt_with_keyring_multithreaded(
40+
n_threads_list: Optional[list] = None,
41+
duration_list: Optional[list] = None,
42+
):
43+
"""Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration."""
44+
# Set defaults if no value is provided
45+
if n_threads_list is None:
46+
n_threads_list = [1, 4, 16, 64]
47+
if duration_list is None:
48+
duration_list = [2, 10, 60]
49+
for n in n_threads_list:
50+
for d in duration_list:
51+
encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=n, duration=d)

0 commit comments

Comments
 (0)