Added max_bytes_encrypted to CMM

seshubaws · Dec 14, 2023 · fe184c4 · fe184c4
1 parent db318cd
commit fe184c4
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 5 deletions.
diff --git a/aws_lambda_powertools/utilities/_data_masking/constants.py b/aws_lambda_powertools/utilities/_data_masking/constants.py
@@ -1,5 +1,9 @@
 DATA_MASKING_STRING: str = "*****"
+# The maximum number of entries that can be retained in the local cryptographic materials cache
 CACHE_CAPACITY: int = 100
+# The maximum time (in seconds) that a cache entry may be kept in the cache
 MAX_CACHE_AGE_SECONDS: float = 300.0
-MAX_MESSAGES_ENCRYPTED: int = 200
-# NOTE: You can also set max messages/bytes per data key
+# Maximum number of messages which are allowed to be encrypted under a single cached data key
+MAX_MESSAGES_ENCRYPTED: int = 4294967296  # 2 ** 32
+# Maximum number of bytes which are allowed to be encrypted under a single cached data key
+MAX_BYTES_ENCRYPTED: int = 9223372036854775807  # 2 ** 63 - 1
diff --git a/aws_lambda_powertools/utilities/_data_masking/provider/kms/aws_encryption_sdk.py b/aws_lambda_powertools/utilities/_data_masking/provider/kms/aws_encryption_sdk.py
@@ -16,6 +16,7 @@
 from aws_lambda_powertools.shared.user_agent import register_feature_to_botocore_session
 from aws_lambda_powertools.utilities._data_masking.constants import (
     CACHE_CAPACITY,
+    MAX_BYTES_ENCRYPTED,
     MAX_CACHE_AGE_SECONDS,
     MAX_MESSAGES_ENCRYPTED,
 )
@@ -67,6 +68,7 @@ def __init__(
         local_cache_capacity: int = CACHE_CAPACITY,
         max_cache_age_seconds: float = MAX_CACHE_AGE_SECONDS,
         max_messages_encrypted: int = MAX_MESSAGES_ENCRYPTED,
+        max_bytes_encrypted: int = MAX_BYTES_ENCRYPTED,
         json_serializer: Callable | None = None,
         json_deserializer: Callable | None = None,
     ):
@@ -77,6 +79,7 @@ def __init__(
             local_cache_capacity=local_cache_capacity,
             max_cache_age_seconds=max_cache_age_seconds,
             max_messages_encrypted=max_messages_encrypted,
+            max_bytes_encrypted=max_bytes_encrypted,
             json_serializer=self.json_serializer,
             json_deserializer=self.json_deserializer,
         )
@@ -103,6 +106,7 @@ def __init__(
         local_cache_capacity: int = CACHE_CAPACITY,
         max_cache_age_seconds: float = MAX_CACHE_AGE_SECONDS,
         max_messages_encrypted: int = MAX_MESSAGES_ENCRYPTED,
+        max_bytes_encrypted: int = MAX_BYTES_ENCRYPTED,
     ):
         session = botocore.session.Session()
         register_feature_to_botocore_session(session, "data-masking")
@@ -118,6 +122,7 @@ def __init__(
             cache=self.cache,
             max_age=max_cache_age_seconds,
             max_messages_encrypted=max_messages_encrypted,
+            max_bytes_encrypted=max_bytes_encrypted,
         )
 
     def encrypt(self, data: bytes | str | Dict | float, **provider_options) -> str:

diff --git a/docs/utilities/data_masking.md b/docs/utilities/data_masking.md
@@ -174,7 +174,8 @@ You have the option to modify some of the configurations we have set as defaults
 | -------------------------- | -------- | ------- | --------------------------------------------------------------------------------------------- |
 | **local_cache_capacity**   |          | `100`   | The maximum number of entries that can be retained in the local cryptographic materials cache |
 | **max_cache_age_seconds**  |          | `300`   | The maximum time (in seconds) that a cache entry may be kept in the cache                     |
-| **max_messages_encrypted** |          | `200`   | The maximum number of messages that may be encrypted under a cache entry                      |
+| **max_messages_encrypted** |          | `4294967296`   | The maximum number of messages that may be encrypted under a cache entry                      |
+| **max_bytes_encrypted** |          | `9223372036854775807`   | The maximum number of bytes that may be encrypted under a cache entry                      |
 
 For more information about the parameters for this provider, please see the [AWS Encryption SDK documentation](https://aws-encryption-sdk-python.readthedocs.io/en/latest/generated/aws_encryption_sdk.materials_managers.caching.html#aws_encryption_sdk.materials_managers.caching.CachingCryptoMaterialsManager){target="_blank" rel="nofollow"}.
 

diff --git a/examples/data_masking/src/getting_started_decrypt_data.py b/examples/data_masking/src/getting_started_decrypt_data.py
@@ -8,7 +8,7 @@
 
 logger = Logger()
 
-KMS_KEY_ARN = os.getenv("KMS_KEY_ARN")
+KMS_KEY_ARN: str = os.getenv("KMS_KEY_ARN")
 encryption_provider = AwsEncryptionSdkProvider(keys=[KMS_KEY_ARN])
 data_masker = DataMasking(provider=encryption_provider)
 

diff --git a/examples/data_masking/src/getting_started_encrypt_data.py b/examples/data_masking/src/getting_started_encrypt_data.py
@@ -8,7 +8,7 @@
 
 logger = Logger()
 
-KMS_KEY_ARN = os.getenv("KMS_KEY_ARN")
+KMS_KEY_ARN: str = os.getenv("KMS_KEY_ARN")
 encryption_provider = AwsEncryptionSdkProvider(keys=[KMS_KEY_ARN])
 data_masker = DataMasking(provider=encryption_provider)