Skip to content

Commit c3fe3e2

Browse files
RitvikKapilaRitvikKapila
committed
chore(v3_default_cmm.py): added test for v3_default_cmm.py; isort fix for migration examples
1 parent 3476816 commit c3fe3e2

6 files changed

+226
-10
lines changed

‎examples/src/legacy/v3_default_cmm.py

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Default crypto material manager class example for ESDK V3."""
4+
import logging
5+
6+
import attr
7+
8+
import aws_encryption_sdk
9+
from aws_encryption_sdk.exceptions import MasterKeyProviderError, SerializationError
10+
from aws_encryption_sdk.identifiers import CommitmentPolicy
11+
from aws_encryption_sdk.internal.crypto.authentication import Signer, Verifier
12+
from aws_encryption_sdk.internal.crypto.elliptic_curve import generate_ecc_signing_key
13+
from aws_encryption_sdk.internal.defaults import ALGORITHM, ALGORITHM_COMMIT_KEY, ENCODED_SIGNER_KEY
14+
from aws_encryption_sdk.internal.str_ops import to_str
15+
from aws_encryption_sdk.internal.utils import prepare_data_keys
16+
from aws_encryption_sdk.internal.utils.commitment import (
17+
validate_commitment_policy_on_decrypt,
18+
validate_commitment_policy_on_encrypt,
19+
)
20+
from aws_encryption_sdk.key_providers.base import MasterKeyProvider
21+
from aws_encryption_sdk.materials_managers import DecryptionMaterials, EncryptionMaterials
22+
from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
23+
24+
_LOGGER = logging.getLogger(__name__)
25+
26+
27+
@attr.s(hash=False)
28+
class V3DefaultCryptoMaterialsManager(CryptoMaterialsManager):
29+
"""Copy of the default crypto material manager for ESDK V3.
30+
31+
This is a copy-paste of the DefaultCryptoMaterialsManager implementation
32+
from the V3 ESDK commit: 98b5eb7c2bd7d1b2a3380aacfa508e8721c4d8a9
33+
This CMM is used to explicitly assert that the V3 implementation of
34+
the DefaultCMM is compatible with future version's logic,
35+
which implicitly asserts that custom implementations of V3-compatible CMMs
36+
are also compatible with future version's logic.
37+
38+
:param master_key_provider: Master key provider to use
39+
:type master_key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
40+
"""
41+
42+
master_key_provider = attr.ib(validator=attr.validators.instance_of(MasterKeyProvider))
43+
44+
def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):
45+
"""Generates a signing key based on the provided algorithm.
46+
47+
:param algorithm: Algorithm for which to generate signing key
48+
:type algorithm: aws_encryption_sdk.identifiers.Algorithm
49+
:param dict encryption_context: Encryption context from request
50+
:returns: Signing key bytes
51+
:rtype: bytes or None
52+
"""
53+
_LOGGER.debug("Generating signing key")
54+
if algorithm.signing_algorithm_info is None:
55+
return None
56+
57+
signer = Signer(algorithm=algorithm, key=generate_ecc_signing_key(algorithm=algorithm))
58+
encryption_context[ENCODED_SIGNER_KEY] = to_str(signer.encoded_public_key())
59+
return signer.key_bytes()
60+
61+
def get_encryption_materials(self, request):
62+
"""Creates encryption materials using underlying master key provider.
63+
64+
:param request: encryption materials request
65+
:type request: aws_encryption_sdk.materials_managers.EncryptionMaterialsRequest
66+
:returns: encryption materials
67+
:rtype: aws_encryption_sdk.materials_managers.EncryptionMaterials
68+
:raises MasterKeyProviderError: if no master keys are available from the underlying master key provider
69+
:raises MasterKeyProviderError: if the primary master key provided by the underlying master key provider
70+
is not included in the full set of master keys provided by that provider
71+
:raises ActionNotAllowedError: if the commitment policy in the request is violated by the algorithm being
72+
used
73+
"""
74+
default_algorithm = ALGORITHM
75+
if request.commitment_policy in (
76+
CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT,
77+
CommitmentPolicy.REQUIRE_ENCRYPT_ALLOW_DECRYPT,
78+
):
79+
default_algorithm = ALGORITHM_COMMIT_KEY
80+
algorithm = request.algorithm if request.algorithm is not None else default_algorithm
81+
82+
validate_commitment_policy_on_encrypt(request.commitment_policy, request.algorithm)
83+
84+
encryption_context = request.encryption_context.copy()
85+
86+
signing_key = self._generate_signing_key_and_update_encryption_context(algorithm, encryption_context)
87+
88+
primary_master_key, master_keys = self.master_key_provider.master_keys_for_encryption(
89+
encryption_context=encryption_context,
90+
plaintext_rostream=request.plaintext_rostream,
91+
plaintext_length=request.plaintext_length,
92+
)
93+
if not master_keys:
94+
raise MasterKeyProviderError("No Master Keys available from Master Key Provider")
95+
if primary_master_key not in master_keys:
96+
raise MasterKeyProviderError("Primary Master Key not in provided Master Keys")
97+
98+
data_encryption_key, encrypted_data_keys = prepare_data_keys(
99+
primary_master_key=primary_master_key,
100+
master_keys=master_keys,
101+
algorithm=algorithm,
102+
encryption_context=encryption_context,
103+
)
104+
105+
_LOGGER.debug("Post-encrypt encryption context: %s", encryption_context)
106+
107+
return EncryptionMaterials(
108+
algorithm=algorithm,
109+
data_encryption_key=data_encryption_key,
110+
encrypted_data_keys=encrypted_data_keys,
111+
encryption_context=encryption_context,
112+
signing_key=signing_key,
113+
)
114+
115+
def _load_verification_key_from_encryption_context(self, algorithm, encryption_context):
116+
"""Loads the verification key from the encryption context if used by algorithm suite.
117+
118+
:param algorithm: Algorithm for which to generate signing key
119+
:type algorithm: aws_encryption_sdk.identifiers.Algorithm
120+
:param dict encryption_context: Encryption context from request
121+
:returns: Raw verification key
122+
:rtype: bytes
123+
:raises SerializationError: if algorithm suite requires message signing and no verification key is found
124+
"""
125+
encoded_verification_key = encryption_context.get(ENCODED_SIGNER_KEY, None)
126+
127+
if algorithm.signing_algorithm_info is not None and encoded_verification_key is None:
128+
raise SerializationError("No signature verification key found in header for signed algorithm.")
129+
130+
if algorithm.signing_algorithm_info is None:
131+
if encoded_verification_key is not None:
132+
raise SerializationError("Signature verification key found in header for non-signed algorithm.")
133+
return None
134+
135+
verifier = Verifier.from_encoded_point(algorithm=algorithm, encoded_point=encoded_verification_key)
136+
return verifier.key_bytes()
137+
138+
def decrypt_materials(self, request):
139+
"""Obtains a plaintext data key from one or more encrypted data keys
140+
using underlying master key provider.
141+
142+
:param request: decrypt materials request
143+
:type request: aws_encryption_sdk.materials_managers.DecryptionMaterialsRequest
144+
:returns: decryption materials
145+
:rtype: aws_encryption_sdk.materials_managers.DecryptionMaterials
146+
"""
147+
validate_commitment_policy_on_decrypt(request.commitment_policy, request.algorithm)
148+
149+
data_key = self.master_key_provider.decrypt_data_key_from_list(
150+
encrypted_data_keys=request.encrypted_data_keys,
151+
algorithm=request.algorithm,
152+
encryption_context=request.encryption_context,
153+
)
154+
verification_key = self._load_verification_key_from_encryption_context(
155+
algorithm=request.algorithm, encryption_context=request.encryption_context
156+
)
157+
158+
return DecryptionMaterials(data_key=data_key, verification_key=verification_key)
159+
160+
161+
def encrypt_decrypt_with_v3_default_cmm(key_arn,
162+
source_plaintext,
163+
botocore_session):
164+
"""Encrypts and decrypts a string using a V3 default CMM.
165+
166+
:param str key_arn: Amazon Resource Name (ARN) of the KMS CMK
167+
:param bytes source_plaintext: Data to encrypt
168+
:param botocore_session: existing botocore session instance
169+
:type botocore_session: botocore.session.Session
170+
"""
171+
# Set up an encryption client with an explicit commitment policy. Note that if you do not explicitly choose a
172+
# commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default.
173+
client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)
174+
175+
# Create a KMS master key provider. Note that because we are planning on decrypting using this same provider,
176+
# we MUST provide the ARN of the KMS Key. If we provide a raw key id or a key alias, decryption will fail.
177+
kms_kwargs = dict(key_ids=[key_arn])
178+
if botocore_session is not None:
179+
kms_kwargs["botocore_session"] = botocore_session
180+
master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
181+
182+
default_cmm = V3DefaultCryptoMaterialsManager(master_key_provider=master_key_provider)
183+
184+
# Encrypt the plaintext source data
185+
ciphertext, encryptor_header = client.encrypt(
186+
source=source_plaintext, materials_manager=default_cmm
187+
)
188+
189+
# Decrypt the ciphertext
190+
cycled_plaintext, decrypted_header = client.decrypt(source=ciphertext, key_provider=master_key_provider)
191+
192+
# Verify that the "cycled" (encrypted, then decrypted) plaintext is identical to the source plaintext
193+
assert cycled_plaintext == source_plaintext
194+
195+
# Verify that the encryption context used in the decrypt operation includes all key pairs from
196+
# the encrypt operation. (The SDK can add pairs, so don't require an exact match.)
197+
#
198+
# In production, always use a meaningful encryption context. In this sample, we omit the
199+
# encryption context (no key pairs).
200+
assert all(
201+
pair in decrypted_header.encryption_context.items() for pair in encryptor_header.encryption_context.items()
202+
)

‎examples/test/legacy/examples_test_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939

4040
from integration_test_utils import ( # noqa pylint: disable=unused-import,import-error
4141
get_cmk_arn,
42-
get_second_cmk_arn,
4342
get_mrk_arn,
43+
get_second_cmk_arn,
4444
get_second_mrk_arn,
4545
)

‎examples/test/legacy/test_i_v3_default_cmm.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Test suite for encryption and decryption using V3 defualt CMM."""
4+
5+
import botocore.session
6+
import pytest
7+
8+
from ...src.legacy.v3_default_cmm import encrypt_decrypt_with_v3_default_cmm
9+
from .examples_test_utils import get_cmk_arn, static_plaintext
10+
11+
pytestmark = [pytest.mark.examples]
12+
13+
14+
def test_v3_default_cmm():
15+
"""Test method for encryption and decryption using V3 defualt CMM"""
16+
plaintext = static_plaintext
17+
cmk_arn = get_cmk_arn()
18+
encrypt_decrypt_with_v3_default_cmm(key_arn=cmk_arn,
19+
source_plaintext=plaintext,
20+
botocore_session=botocore.session.Session())

‎examples/test/migration/test_i_migration_aws_kms_key_example.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@
33
"""Test suite for the migration_aws_kms_key_example."""
44
import pytest
55

6-
from ...src.migration.migration_aws_kms_key_example import (
7-
migration_aws_kms_key,
8-
)
6+
from ...src.migration.migration_aws_kms_key_example import migration_aws_kms_key
97

108
pytestmark = [pytest.mark.examples]
119

‎examples/test/migration/test_i_migration_raw_aes_key_example.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@
33
"""Test suite for the migration_raw_aes_key_example."""
44
import pytest
55

6-
from ...src.migration.migration_raw_aes_key_example import (
7-
migration_raw_aes_key,
8-
)
6+
from ...src.migration.migration_raw_aes_key_example import migration_raw_aes_key
97

108
pytestmark = [pytest.mark.examples]
119

‎examples/test/migration/test_i_migration_raw_rsa_key_example.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@
33
"""Test suite for the migration_raw_rsa_key_example."""
44
import pytest
55

6-
from ...src.migration.migration_raw_rsa_key_example import (
7-
migration_raw_rsa_key,
8-
)
6+
from ...src.migration.migration_raw_rsa_key_example import migration_raw_rsa_key
97

108
pytestmark = [pytest.mark.examples]
119

0 commit comments

Comments
 (0)