From ad1d90c5410fcceb3907ba2477905b252e3bef24 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 11:37:47 -0400 Subject: [PATCH 01/16] Started --- datadog_lambda/metric.py | 43 +++++++++++++++++++++++++++---- tests/test_metric.py | 55 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 5eb1c2ac..e4e06345 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -9,13 +9,14 @@ import base64 import logging +from botocore.exceptions import ClientError import boto3 from datadog import api, initialize, statsd from datadog.threadstats import ThreadStats from datadog_lambda.extension import should_use_extension from datadog_lambda.tags import get_enhanced_metrics_tags, tag_dd_lambda_layer - +KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" ENHANCED_METRICS_NAMESPACE_PREFIX = "aws.lambda.enhanced" logger = logging.getLogger(__name__) @@ -212,13 +213,45 @@ def submit_errors_metric(lambda_context): """ submit_enhanced_metric("errors", lambda_context) +def decrypt_kms_api_key(ciphertext): + """ + Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS. + For this to work properly, the Lambda function must have the appropriate IAM permissions. + + Args: + ciphertext (string): The base64-encoded ciphertext to decrypt + """ + decoded_bytes = base64.b64decode(ciphertext) + + """ + The Lambda console UI changed the way it encrypts environment variables. + The current behavior as of May 2021 is to encrypt environment variables using the function name as an encryption context. + Previously, the behavior was to encrypt environment variables without an encryption context. + We need to try both, as supplying the incorrect encryption context will cause decryption to fail. + """ + # Try with encryption context + function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME") + try: + plaintext = kms_client.decrypt( + CiphertextBlob=decoded_bytes, + EncryptionContext={ + KMS_ENCRYPTION_CONTEXT_KEY: function_name, + } + )["Plaintext"] + except ClientError: + logger.debug("Failed to decrypt ciphertext with encryption context, retrying without encryption context") + # Try without encryption context + plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)["Plaintext"] + + return plaintext -# Set API Key and Host in the module, so they only set once per container +# Set API Key if not api._api_key: DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "") DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "") DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "") DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")) + if DD_API_KEY_SECRET_ARN: api._api_key = boto3.client("secretsmanager").get_secret_value( SecretId=DD_API_KEY_SECRET_ARN @@ -228,11 +261,11 @@ def submit_errors_metric(lambda_context): Name=DD_API_KEY_SSM_NAME, WithDecryption=True )["Parameter"]["Value"] elif DD_KMS_API_KEY: - api._api_key = boto3.client("kms").decrypt( - CiphertextBlob=base64.b64decode(DD_KMS_API_KEY) - )["Plaintext"] + kms_client = boto3.client("kms") + api._api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY) else: api._api_key = DD_API_KEY + logger.debug("Setting DATADOG_API_KEY of length %d", len(api._api_key)) # Set DATADOG_HOST, to send data to a non-default Datadog datacenter diff --git a/tests/test_metric.py b/tests/test_metric.py index 848a9328..5c13a4cf 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -6,8 +6,14 @@ except ImportError: from mock import patch, call +from botocore.exceptions import ClientError as BotocoreClientError from datadog.api.exceptions import ClientError -from datadog_lambda.metric import lambda_metric, ThreadStatsWriter +from datadog_lambda.metric import ( + decrypt_kms_api_key, + lambda_metric, + ThreadStatsWriter, + KMS_ENCRYPTION_CONTEXT_KEY +) from datadog_lambda.tags import _format_dd_lambda_layer_tag @@ -58,3 +64,50 @@ def test_retry_on_remote_disconnected(self): ) lambda_stats.flush() self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2) + + +class TestDecryptKMSApiKey(unittest.TestCase): + + mock_function_name = "myFunction" + + # An API key encrypted with KMS and encoded as a base64 string + mock_encrypted_api_key_base64 = "MjIyMjIyMjIyMjIyMjIyMg==" + + # The encrypted API key after it has been decoded from base64 + mock_encrypted_api_key = "2222222222222222" + + # The true value of the API key after decryption by KMS + expected_decrypted_api_key = "1111111111111111" + + def test_key_encrypted_with_encryption_context(self): + os.environ["AWS_LAMBDA_FUNCTION_NAME"] = self.mock_function_name + + class MockKMSClient: + def decrypt(CiphertextBlob, EncryptionContext): + if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != self.mock_function_name: + raise BotocoreClientError() + if CiphertextBlob == self.mock_encrypted_api_key.encode('utf-8'): + return { + "Plaintext": self.expected_decrypted_api_key, + } + + mock_kms_client = MockKMSClient() + decrypted_key = decrypt_kms_api_key(mock_kms_client, self.mock_encrypted_api_key_base64) + self.assertEqual(decrypted_key, self.expected_decrypted_api_key) + + del os.environ["AWS_LAMBDA_FUNCTION_NAME"] + + def test_key_encrypted_without_encryption_context(self): + + class MockKMSClient: + def decrypt(CiphertextBlob, EncryptionContext): + if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: + raise BotocoreClientError() + if CiphertextBlob == self.mock_encrypted_api_key.encode('utf-8'): + return { + "Plaintext": self.expected_decrypted_api_key, + } + + mock_kms_client = MockKMSClient() + decrypted_key = decrypt_kms_api_key(mock_kms_client, self.mock_encrypted_api_key_base64) + self.assertEqual(decrypted_key, self.expected_decrypted_api_key) From 3991df7033581904c07377fb8a3f1b6a14520917 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 11:40:57 -0400 Subject: [PATCH 02/16] Include additional file, black --- datadog_lambda/cold_start.py | 6 ++---- datadog_lambda/metric.py | 11 ++++++++--- datadog_lambda/module_name.py | 3 +-- datadog_lambda/tags.py | 12 ++++-------- datadog_lambda/tracing.py | 14 ++++++++++---- 5 files changed, 25 insertions(+), 21 deletions(-) diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py index e53341f4..c8862bf1 100644 --- a/datadog_lambda/cold_start.py +++ b/datadog_lambda/cold_start.py @@ -14,12 +14,10 @@ def set_cold_start(): def is_cold_start(): - """Returns the value of the global cold_start - """ + """Returns the value of the global cold_start""" return _cold_start def get_cold_start_tag(): - """Returns the cold start tag to be used in metrics - """ + """Returns the cold start tag to be used in metrics""" return "cold_start:{}".format(str(is_cold_start()).lower()) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index e4e06345..7e4172fd 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -213,12 +213,14 @@ def submit_errors_metric(lambda_context): """ submit_enhanced_metric("errors", lambda_context) -def decrypt_kms_api_key(ciphertext): + +def decrypt_kms_api_key(kms_client, ciphertext): """ Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS. For this to work properly, the Lambda function must have the appropriate IAM permissions. Args: + kms_client: The KMS client to use for decryption ciphertext (string): The base64-encoded ciphertext to decrypt """ decoded_bytes = base64.b64decode(ciphertext) @@ -236,15 +238,18 @@ def decrypt_kms_api_key(ciphertext): CiphertextBlob=decoded_bytes, EncryptionContext={ KMS_ENCRYPTION_CONTEXT_KEY: function_name, - } + }, )["Plaintext"] except ClientError: - logger.debug("Failed to decrypt ciphertext with encryption context, retrying without encryption context") + logger.debug( + "Failed to decrypt ciphertext with encryption context, retrying without encryption context" + ) # Try without encryption context plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)["Plaintext"] return plaintext + # Set API Key if not api._api_key: DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "") diff --git a/datadog_lambda/module_name.py b/datadog_lambda/module_name.py index f018c5d6..9e4a93e5 100644 --- a/datadog_lambda/module_name.py +++ b/datadog_lambda/module_name.py @@ -1,4 +1,3 @@ def modify_module_name(module_name): - """Returns a valid modified module to get imported - """ + """Returns a valid modified module to get imported""" return ".".join(module_name.split("/")) diff --git a/datadog_lambda/tags.py b/datadog_lambda/tags.py index f21c8f92..a19a32ad 100644 --- a/datadog_lambda/tags.py +++ b/datadog_lambda/tags.py @@ -69,8 +69,7 @@ def parse_lambda_tags_from_arn(lambda_context): def get_runtime_tag(): - """Get the runtime tag from the current Python version - """ + """Get the runtime tag from the current Python version""" major_version, minor_version, _ = python_version_tuple() return "runtime:python{major}.{minor}".format( @@ -79,14 +78,12 @@ def get_runtime_tag(): def get_library_version_tag(): - """Get datadog lambda library tag - """ + """Get datadog lambda library tag""" return "datadog_lambda:v{}".format(__version__) def get_enhanced_metrics_tags(lambda_context): - """Get the list of tags to apply to enhanced metrics - """ + """Get the list of tags to apply to enhanced metrics""" return parse_lambda_tags_from_arn(lambda_context) + [ get_cold_start_tag(), "memorysize:{}".format(lambda_context.memory_limit_in_mb), @@ -96,8 +93,7 @@ def get_enhanced_metrics_tags(lambda_context): def check_if_number(alias): - """ Check if the alias is a version or number. Python 2 has no easy way to test this like Python 3 - """ + """Check if the alias is a version or number. Python 2 has no easy way to test this like Python 3""" try: float(alias) return True diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 4f528cc3..ae48a161 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -186,7 +186,11 @@ def extract_context_custom_extractor(extractor, event, lambda_context): Extract Datadog trace context using a custom trace extractor function """ try: - (trace_id, parent_id, sampling_priority,) = extractor(event, lambda_context) + ( + trace_id, + parent_id, + sampling_priority, + ) = extractor(event, lambda_context) return trace_id, parent_id, sampling_priority except Exception as e: logger.debug("The trace extractor returned with error %s", e) @@ -205,9 +209,11 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): trace_context_source = None if extractor is not None: - (trace_id, parent_id, sampling_priority,) = extract_context_custom_extractor( - extractor, event, lambda_context - ) + ( + trace_id, + parent_id, + sampling_priority, + ) = extract_context_custom_extractor(extractor, event, lambda_context) elif "headers" in event: ( trace_id, From 5e7e1841c323cb250a9d9c2dbbfd3f331137d6bf Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 14:12:02 -0400 Subject: [PATCH 03/16] Fix kwargs, black version --- scripts/check_format.sh | 2 +- tests/test_metric.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/check_format.sh b/scripts/check_format.sh index 7c38d48b..0c72125b 100755 --- a/scripts/check_format.sh +++ b/scripts/check_format.sh @@ -6,7 +6,7 @@ if [ "$PYTHON_VERSION" = "2" ]; then echo "Skipping formatting, black not compatible with python 2" exit 0 fi -pip install -Iv black==19.10b0 +pip install -Iv black==20.8b1 python -m black --check datadog_lambda/ --diff python -m black --check tests --diff diff --git a/tests/test_metric.py b/tests/test_metric.py index 5c13a4cf..3c828cfb 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -83,7 +83,7 @@ def test_key_encrypted_with_encryption_context(self): os.environ["AWS_LAMBDA_FUNCTION_NAME"] = self.mock_function_name class MockKMSClient: - def decrypt(CiphertextBlob, EncryptionContext): + def decrypt(CiphertextBlob=None, EncryptionContext=None): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != self.mock_function_name: raise BotocoreClientError() if CiphertextBlob == self.mock_encrypted_api_key.encode('utf-8'): @@ -100,7 +100,7 @@ def decrypt(CiphertextBlob, EncryptionContext): def test_key_encrypted_without_encryption_context(self): class MockKMSClient: - def decrypt(CiphertextBlob, EncryptionContext): + def decrypt(CiphertextBlob=None, EncryptionContext=None): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: raise BotocoreClientError() if CiphertextBlob == self.mock_encrypted_api_key.encode('utf-8'): From 34a15041da4749123124cac8f3bad55dd5e1a1a3 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 14:16:57 -0400 Subject: [PATCH 04/16] Black on tests --- tests/test_metric.py | 26 ++++++++++++++++---------- tests/test_tracing.py | 37 ++++++++++++++++++++++++++++++------- 2 files changed, 46 insertions(+), 17 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index 3c828cfb..a00581f1 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -12,7 +12,7 @@ decrypt_kms_api_key, lambda_metric, ThreadStatsWriter, - KMS_ENCRYPTION_CONTEXT_KEY + KMS_ENCRYPTION_CONTEXT_KEY, ) from datadog_lambda.tags import _format_dd_lambda_layer_tag @@ -84,30 +84,36 @@ def test_key_encrypted_with_encryption_context(self): class MockKMSClient: def decrypt(CiphertextBlob=None, EncryptionContext=None): - if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != self.mock_function_name: + if ( + EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) + != self.mock_function_name + ): raise BotocoreClientError() - if CiphertextBlob == self.mock_encrypted_api_key.encode('utf-8'): + if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): return { "Plaintext": self.expected_decrypted_api_key, - } + } mock_kms_client = MockKMSClient() - decrypted_key = decrypt_kms_api_key(mock_kms_client, self.mock_encrypted_api_key_base64) + decrypted_key = decrypt_kms_api_key( + mock_kms_client, self.mock_encrypted_api_key_base64 + ) self.assertEqual(decrypted_key, self.expected_decrypted_api_key) del os.environ["AWS_LAMBDA_FUNCTION_NAME"] - - def test_key_encrypted_without_encryption_context(self): + def test_key_encrypted_without_encryption_context(self): class MockKMSClient: def decrypt(CiphertextBlob=None, EncryptionContext=None): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: raise BotocoreClientError() - if CiphertextBlob == self.mock_encrypted_api_key.encode('utf-8'): + if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): return { "Plaintext": self.expected_decrypted_api_key, } - + mock_kms_client = MockKMSClient() - decrypted_key = decrypt_kms_api_key(mock_kms_client, self.mock_encrypted_api_key_base64) + decrypted_key = decrypt_kms_api_key( + mock_kms_client, self.mock_encrypted_api_key_base64 + ) self.assertEqual(decrypted_key, self.expected_decrypted_api_key) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index b7365df4..431f339f 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -73,7 +73,8 @@ def test_without_datadog_trace_headers(self): ctx, source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "xray") self.assertDictEqual( - ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, + ctx, + {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -93,7 +94,8 @@ def test_with_incomplete_datadog_trace_headers(self): ) self.assertEqual(source, "xray") self.assertDictEqual( - ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, + ctx, + {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -118,7 +120,8 @@ def test_with_complete_datadog_trace_headers(self): ) self.assertEqual(source, "event") self.assertDictEqual( - ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, + ctx, + {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, ) self.assertDictEqual( get_dd_trace_context(), @@ -160,7 +163,12 @@ def extractor_foo(event, context): ) self.assertEquals(ctx_source, "event") self.assertDictEqual( - ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",}, + ctx, + { + "trace-id": "123", + "parent-id": "321", + "sampling-priority": "1", + }, ) self.assertDictEqual( get_dd_trace_context(), @@ -189,7 +197,12 @@ def extractor_raiser(event, context): ) self.assertEquals(ctx_source, "xray") self.assertDictEqual( - ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2",}, + ctx, + { + "trace-id": "4369", + "parent-id": "65535", + "sampling-priority": "2", + }, ) self.assertDictEqual( get_dd_trace_context(), @@ -235,7 +248,12 @@ def test_with_sqs_distributed_datadog_trace_data(self): ctx, source = extract_dd_trace_context(sqs_event, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( - ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",}, + ctx, + { + "trace-id": "123", + "parent-id": "321", + "sampling-priority": "1", + }, ) self.assertDictEqual( get_dd_trace_context(), @@ -267,7 +285,12 @@ def test_with_client_context_datadog_trace_data(self): ctx, source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( - ctx, {"trace-id": "666", "parent-id": "777", "sampling-priority": "1",}, + ctx, + { + "trace-id": "666", + "parent-id": "777", + "sampling-priority": "1", + }, ) self.assertDictEqual( get_dd_trace_context(), From 4a66198528bc7f8124e05af9935da53ec28d0cce Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 14:21:19 -0400 Subject: [PATCH 05/16] Fix mocks --- tests/test_metric.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index a00581f1..47635451 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -83,7 +83,7 @@ def test_key_encrypted_with_encryption_context(self): os.environ["AWS_LAMBDA_FUNCTION_NAME"] = self.mock_function_name class MockKMSClient: - def decrypt(CiphertextBlob=None, EncryptionContext=None): + def decrypt(self, CiphertextBlob, EncryptionContext): if ( EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != self.mock_function_name @@ -104,7 +104,7 @@ def decrypt(CiphertextBlob=None, EncryptionContext=None): def test_key_encrypted_without_encryption_context(self): class MockKMSClient: - def decrypt(CiphertextBlob=None, EncryptionContext=None): + def decrypt(self, CiphertextBlob, EncryptionContext): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: raise BotocoreClientError() if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): From 29a64c47792b53447ad49aafb1bd86a23dd8d83b Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 14:24:50 -0400 Subject: [PATCH 06/16] Fix comments --- datadog_lambda/metric.py | 9 +++++---- datadog_lambda/tags.py | 5 ++++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 7e4172fd..7961960e 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -226,10 +226,11 @@ def decrypt_kms_api_key(kms_client, ciphertext): decoded_bytes = base64.b64decode(ciphertext) """ - The Lambda console UI changed the way it encrypts environment variables. - The current behavior as of May 2021 is to encrypt environment variables using the function name as an encryption context. - Previously, the behavior was to encrypt environment variables without an encryption context. - We need to try both, as supplying the incorrect encryption context will cause decryption to fail. + The Lambda console UI changed the way it encrypts environment variables. The current behavior + as of May 2021 is to encrypt environment variables using the function name as an encryption + context. Previously, the behavior was to encrypt environment variables without an encryption + context. We need to try both, as supplying the incorrect encryption context will cause + decryption to fail. """ # Try with encryption context function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME") diff --git a/datadog_lambda/tags.py b/datadog_lambda/tags.py index a19a32ad..2ba6ae72 100644 --- a/datadog_lambda/tags.py +++ b/datadog_lambda/tags.py @@ -93,7 +93,10 @@ def get_enhanced_metrics_tags(lambda_context): def check_if_number(alias): - """Check if the alias is a version or number. Python 2 has no easy way to test this like Python 3""" + """ + Check if the alias is a version or number. + Python 2 has no easy way to test this like Python 3 + """ try: float(alias) return True From 5e4671fcfe9033269e392e0ec082561bfda2ad17 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 14:28:55 -0400 Subject: [PATCH 07/16] Black 21 --- scripts/check_format.sh | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/check_format.sh b/scripts/check_format.sh index 0c72125b..d616017e 100755 --- a/scripts/check_format.sh +++ b/scripts/check_format.sh @@ -6,7 +6,7 @@ if [ "$PYTHON_VERSION" = "2" ]; then echo "Skipping formatting, black not compatible with python 2" exit 0 fi -pip install -Iv black==20.8b1 +pip install -Iv black==21.5b2 python -m black --check datadog_lambda/ --diff python -m black --check tests --diff diff --git a/setup.py b/setup.py index 921aab4c..820f320f 100644 --- a/setup.py +++ b/setup.py @@ -35,7 +35,7 @@ # If building for Python 3, use the latest version of setuptools "setuptools>=54.2.0; python_version >= '3.0'", # If building for Python 2, use the latest version that supports Python 2 - "setuptools>=44.1.1; python_version < '3.0'" + "setuptools>=44.1.1; python_version < '3.0'", ], extras_require={ "dev": ["nose2==0.9.1", "flake8==3.7.9", "requests==2.22.0", "boto3==1.10.33"] From 203791ee2670177a60e902f662b7cbbbd5f357f0 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 15:28:50 -0400 Subject: [PATCH 08/16] More fixes --- datadog_lambda/metric.py | 2 +- tests/test_metric.py | 11 +++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 7961960e..15dfb753 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -243,7 +243,7 @@ def decrypt_kms_api_key(kms_client, ciphertext): )["Plaintext"] except ClientError: logger.debug( - "Failed to decrypt ciphertext with encryption context, retrying without encryption context" + "Failed to decrypt ciphertext with encryption context, retrying without" ) # Try without encryption context plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)["Plaintext"] diff --git a/tests/test_metric.py b/tests/test_metric.py index 47635451..cd12ff7d 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -16,6 +16,7 @@ ) from datadog_lambda.tags import _format_dd_lambda_layer_tag +MOCK_FUNCTION_NAME = "myFunction" class TestLambdaMetric(unittest.TestCase): def setUp(self): @@ -68,8 +69,6 @@ def test_retry_on_remote_disconnected(self): class TestDecryptKMSApiKey(unittest.TestCase): - mock_function_name = "myFunction" - # An API key encrypted with KMS and encoded as a base64 string mock_encrypted_api_key_base64 = "MjIyMjIyMjIyMjIyMjIyMg==" @@ -80,15 +79,15 @@ class TestDecryptKMSApiKey(unittest.TestCase): expected_decrypted_api_key = "1111111111111111" def test_key_encrypted_with_encryption_context(self): - os.environ["AWS_LAMBDA_FUNCTION_NAME"] = self.mock_function_name + os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME class MockKMSClient: def decrypt(self, CiphertextBlob, EncryptionContext): if ( EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) - != self.mock_function_name + != MOCK_FUNCTION_NAME ): - raise BotocoreClientError() + raise BotocoreClientError("Error", "Decrypt") if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): return { "Plaintext": self.expected_decrypted_api_key, @@ -106,7 +105,7 @@ def test_key_encrypted_without_encryption_context(self): class MockKMSClient: def decrypt(self, CiphertextBlob, EncryptionContext): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: - raise BotocoreClientError() + raise BotocoreClientError("Error", "Decrypt") if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): return { "Plaintext": self.expected_decrypted_api_key, From e02e11d76ae43f5d5e23ab396b7281757c258390 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 15:30:28 -0400 Subject: [PATCH 09/16] Black --- tests/test_metric.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_metric.py b/tests/test_metric.py index cd12ff7d..caf9df5c 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -18,6 +18,7 @@ MOCK_FUNCTION_NAME = "myFunction" + class TestLambdaMetric(unittest.TestCase): def setUp(self): patcher = patch("datadog_lambda.metric.lambda_stats") From 608eb26f404f96f2c14e267d36d4e2820a5152dd Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 15:51:45 -0400 Subject: [PATCH 10/16] Botocore error --- tests/test_metric.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index caf9df5c..d142e31c 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -88,7 +88,7 @@ def decrypt(self, CiphertextBlob, EncryptionContext): EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != MOCK_FUNCTION_NAME ): - raise BotocoreClientError("Error", "Decrypt") + raise BotocoreClientError({}, "Decrypt") if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): return { "Plaintext": self.expected_decrypted_api_key, @@ -106,7 +106,7 @@ def test_key_encrypted_without_encryption_context(self): class MockKMSClient: def decrypt(self, CiphertextBlob, EncryptionContext): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: - raise BotocoreClientError("Error", "Decrypt") + raise BotocoreClientError({}, "Decrypt") if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): return { "Plaintext": self.expected_decrypted_api_key, From b62eba2ac33f4b3d71d0591d6959e26e89f29dc6 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 15:54:44 -0400 Subject: [PATCH 11/16] Positional arguments --- tests/test_metric.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index d142e31c..9032912e 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -83,7 +83,7 @@ def test_key_encrypted_with_encryption_context(self): os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME class MockKMSClient: - def decrypt(self, CiphertextBlob, EncryptionContext): + def decrypt(self, CiphertextBlob=None, EncryptionContext=None): if ( EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != MOCK_FUNCTION_NAME @@ -104,7 +104,7 @@ def decrypt(self, CiphertextBlob, EncryptionContext): def test_key_encrypted_without_encryption_context(self): class MockKMSClient: - def decrypt(self, CiphertextBlob, EncryptionContext): + def decrypt(self, CiphertextBlob=None, EncryptionContext=None): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: raise BotocoreClientError({}, "Decrypt") if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): From 311e49c26b37cb678fb4dc3474cf7dbb6aef0b47 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 16:00:19 -0400 Subject: [PATCH 12/16] Fix again --- tests/test_metric.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index 9032912e..1bae52c5 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -83,7 +83,7 @@ def test_key_encrypted_with_encryption_context(self): os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME class MockKMSClient: - def decrypt(self, CiphertextBlob=None, EncryptionContext=None): + def decrypt(self, CiphertextBlob=None, EncryptionContext={}): if ( EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != MOCK_FUNCTION_NAME @@ -104,7 +104,7 @@ def decrypt(self, CiphertextBlob=None, EncryptionContext=None): def test_key_encrypted_without_encryption_context(self): class MockKMSClient: - def decrypt(self, CiphertextBlob=None, EncryptionContext=None): + def decrypt(self, CiphertextBlob=None, EncryptionContext={}): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: raise BotocoreClientError({}, "Decrypt") if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): From 51ab150f0783db5601c86097c2fe6541b48b379d Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 17:06:56 -0400 Subject: [PATCH 13/16] Move stuff out of class --- tests/test_metric.py | 35 ++++++++++++++++------------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index 1bae52c5..9017d26b 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -16,9 +16,6 @@ ) from datadog_lambda.tags import _format_dd_lambda_layer_tag -MOCK_FUNCTION_NAME = "myFunction" - - class TestLambdaMetric(unittest.TestCase): def setUp(self): patcher = patch("datadog_lambda.metric.lambda_stats") @@ -67,17 +64,17 @@ def test_retry_on_remote_disconnected(self): lambda_stats.flush() self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2) +MOCK_FUNCTION_NAME = "myFunction" -class TestDecryptKMSApiKey(unittest.TestCase): - - # An API key encrypted with KMS and encoded as a base64 string - mock_encrypted_api_key_base64 = "MjIyMjIyMjIyMjIyMjIyMg==" +# An API key encrypted with KMS and encoded as a base64 string +MOCK_ENCRYPTED_API_KEY_BASE64 = "MjIyMjIyMjIyMjIyMjIyMg==" - # The encrypted API key after it has been decoded from base64 - mock_encrypted_api_key = "2222222222222222" +# The encrypted API key after it has been decoded from base64 +MOCK_ENCRYPTED_API_KEY = "2222222222222222" - # The true value of the API key after decryption by KMS - expected_decrypted_api_key = "1111111111111111" +# The true value of the API key after decryption by KMS +EXPECTED_DECRYPTED_API_KEY = "1111111111111111" +class TestDecryptKMSApiKey(unittest.TestCase): def test_key_encrypted_with_encryption_context(self): os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME @@ -89,16 +86,16 @@ def decrypt(self, CiphertextBlob=None, EncryptionContext={}): != MOCK_FUNCTION_NAME ): raise BotocoreClientError({}, "Decrypt") - if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): + if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"): return { - "Plaintext": self.expected_decrypted_api_key, + "Plaintext": EXPECTED_DECRYPTED_API_KEY, } mock_kms_client = MockKMSClient() decrypted_key = decrypt_kms_api_key( - mock_kms_client, self.mock_encrypted_api_key_base64 + mock_kms_client, MOCK_ENCRYPTED_API_KEY_BASE64 ) - self.assertEqual(decrypted_key, self.expected_decrypted_api_key) + self.assertEqual(decrypted_key, EXPECTED_DECRYPTED_API_KEY) del os.environ["AWS_LAMBDA_FUNCTION_NAME"] @@ -107,13 +104,13 @@ class MockKMSClient: def decrypt(self, CiphertextBlob=None, EncryptionContext={}): if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None: raise BotocoreClientError({}, "Decrypt") - if CiphertextBlob == self.mock_encrypted_api_key.encode("utf-8"): + if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"): return { - "Plaintext": self.expected_decrypted_api_key, + "Plaintext": EXPECTED_DECRYPTED_API_KEY, } mock_kms_client = MockKMSClient() decrypted_key = decrypt_kms_api_key( - mock_kms_client, self.mock_encrypted_api_key_base64 + mock_kms_client, MOCK_ENCRYPTED_API_KEY_BASE64 ) - self.assertEqual(decrypted_key, self.expected_decrypted_api_key) + self.assertEqual(decrypted_key, EXPECTED_DECRYPTED_API_KEY) From 267fefeb65b88e971b6b36f1641b0fa03af1dce8 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Wed, 9 Jun 2021 17:09:05 -0400 Subject: [PATCH 14/16] Black --- tests/test_metric.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index 9017d26b..799acf82 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -16,6 +16,7 @@ ) from datadog_lambda.tags import _format_dd_lambda_layer_tag + class TestLambdaMetric(unittest.TestCase): def setUp(self): patcher = patch("datadog_lambda.metric.lambda_stats") @@ -64,6 +65,7 @@ def test_retry_on_remote_disconnected(self): lambda_stats.flush() self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2) + MOCK_FUNCTION_NAME = "myFunction" # An API key encrypted with KMS and encoded as a base64 string @@ -74,8 +76,9 @@ def test_retry_on_remote_disconnected(self): # The true value of the API key after decryption by KMS EXPECTED_DECRYPTED_API_KEY = "1111111111111111" -class TestDecryptKMSApiKey(unittest.TestCase): + +class TestDecryptKMSApiKey(unittest.TestCase): def test_key_encrypted_with_encryption_context(self): os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME From 63a1a4777ca2eaefc8df7854b56d017da2410a01 Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Thu, 10 Jun 2021 11:22:28 -0400 Subject: [PATCH 15/16] Encoding for Python 3 --- datadog_lambda/metric.py | 4 ++-- tests/test_metric.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 15dfb753..16703a38 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -240,13 +240,13 @@ def decrypt_kms_api_key(kms_client, ciphertext): EncryptionContext={ KMS_ENCRYPTION_CONTEXT_KEY: function_name, }, - )["Plaintext"] + )["Plaintext"].decode('utf-8') except ClientError: logger.debug( "Failed to decrypt ciphertext with encryption context, retrying without" ) # Try without encryption context - plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)["Plaintext"] + plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)["Plaintext"].decode('utf-8') return plaintext diff --git a/tests/test_metric.py b/tests/test_metric.py index 799acf82..d22cf881 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -91,7 +91,7 @@ def decrypt(self, CiphertextBlob=None, EncryptionContext={}): raise BotocoreClientError({}, "Decrypt") if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"): return { - "Plaintext": EXPECTED_DECRYPTED_API_KEY, + "Plaintext": EXPECTED_DECRYPTED_API_KEY.encode("utf-8"), } mock_kms_client = MockKMSClient() @@ -109,7 +109,7 @@ def decrypt(self, CiphertextBlob=None, EncryptionContext={}): raise BotocoreClientError({}, "Decrypt") if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"): return { - "Plaintext": EXPECTED_DECRYPTED_API_KEY, + "Plaintext": EXPECTED_DECRYPTED_API_KEY.encode("utf-8"), } mock_kms_client = MockKMSClient() From 9c22c28382f2105258001b20465e7602b76fa7cb Mon Sep 17 00:00:00 2001 From: Nick Hinsch Date: Thu, 10 Jun 2021 11:25:43 -0400 Subject: [PATCH 16/16] Black --- datadog_lambda/metric.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 16703a38..62100464 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -240,13 +240,15 @@ def decrypt_kms_api_key(kms_client, ciphertext): EncryptionContext={ KMS_ENCRYPTION_CONTEXT_KEY: function_name, }, - )["Plaintext"].decode('utf-8') + )["Plaintext"].decode("utf-8") except ClientError: logger.debug( "Failed to decrypt ciphertext with encryption context, retrying without" ) # Try without encryption context - plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)["Plaintext"].decode('utf-8') + plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)[ + "Plaintext" + ].decode("utf-8") return plaintext