diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index c539ea05..fd3e2c17 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -1,6 +1,5 @@ import os import logging -import base64 logger = logging.getLogger(__name__) KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" @@ -9,6 +8,7 @@ def decrypt_kms_api_key(kms_client, ciphertext): from botocore.exceptions import ClientError + import base64 """ Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS. diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 0fae76dd..9a27673c 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -2,10 +2,8 @@ # under the Apache License Version 2.0. # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import hashlib import logging import os -import base64 import traceback import ujson as json from datetime import datetime, timezone @@ -259,6 +257,8 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): dd_json_data = None dd_json_data_type = dd_payload.get("Type") or dd_payload.get("dataType") if dd_json_data_type == "Binary": + import base64 + dd_json_data = dd_payload.get("binaryValue") or dd_payload.get("Value") if dd_json_data: dd_json_data = base64.b64decode(dd_json_data) @@ -373,6 +373,8 @@ def extract_context_from_kinesis_event(event, lambda_context): return extract_context_from_lambda_context(lambda_context) data = kinesis.get("data") if data: + import base64 + b64_bytes = data.encode("ascii") str_bytes = base64.b64decode(b64_bytes) data_str = str_bytes.decode("ascii") @@ -387,6 +389,8 @@ def extract_context_from_kinesis_event(event, lambda_context): def _deterministic_sha256_hash(s: str, part: str) -> int: + import hashlib + sha256_hash = hashlib.sha256(s.encode()).hexdigest() # First two chars is '0b'. zfill to ensure 256 bits, but we only care about the first 128 bits binary_hash = bin(int(sha256_hash, 16))[2:].zfill(256) @@ -551,6 +555,8 @@ def get_injected_authorizer_data(event, is_http_api) -> dict: if not dd_data_raw: return None + import base64 + injected_data = json.loads(base64.b64decode(dd_data_raw)) # Lambda authorizer's results can be cached. But the payload will still have the injected diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index 708138bf..8090e36e 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -3,7 +3,6 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import base64 import gzip import ujson as json from io import BytesIO, BufferedReader @@ -242,6 +241,8 @@ def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> s # e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz if source.event_type == EventTypes.CLOUDWATCH_LOGS: + import base64 + with gzip.GzipFile( fileobj=BytesIO(base64.b64decode(event.get("awslogs", {}).get("data"))) ) as decompress_stream: diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 5641bd15..8c1914e3 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0. # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import base64 import os import logging import traceback @@ -23,11 +22,6 @@ XraySubsegment, Headers, ) -from datadog_lambda.metric import ( - flush_stats, - submit_invocations_metric, - submit_errors_metric, -) from datadog_lambda.module_name import modify_module_name from datadog_lambda.patch import patch_all from datadog_lambda.span_pointers import calculate_span_pointers @@ -248,7 +242,11 @@ def __call__(self, event, context, **kwargs): self.response = self.func(event, context, **kwargs) return self.response except Exception: - submit_errors_metric(context) + if not should_use_extension: + from datadog_lambda.metric import submit_errors_metric + + submit_errors_metric(context) + if self.span: self.span.set_traceback() raise @@ -274,6 +272,9 @@ def _inject_authorizer_span_headers(self, request_id): injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns if request_id is not None: injected_headers[Headers.Authorizing_Request_Id] = request_id + + import base64 + datadog_data = base64.b64encode( json.dumps(injected_headers, escape_forward_slashes=False).encode() ).decode() @@ -284,7 +285,12 @@ def _before(self, event, context): try: self.response = None set_cold_start(init_timestamp_ns) - submit_invocations_metric(context) + + if not should_use_extension: + from datadog_lambda.metric import submit_invocations_metric + + submit_invocations_metric(context) + self.trigger_tags = extract_trigger_tags(event, context) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( @@ -383,6 +389,8 @@ def _after(self, event, context): logger.debug("Failed to create cold start spans. %s", e) if not self.flush_to_log or should_use_extension: + from datadog_lambda.metric import flush_stats + flush_stats(context) if should_use_extension and self.local_testing_mode: # when testing locally, the extension does not know when an diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index f47285e6..4b243036 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -470,7 +470,7 @@ def lambda_handler(event, context): self.mock_write_metric_point_to_stdout.assert_not_called() def test_only_one_wrapper_in_use(self): - patcher = patch("datadog_lambda.wrapper.submit_invocations_metric") + patcher = patch("datadog_lambda.metric.submit_invocations_metric") self.mock_submit_invocations_metric = patcher.start() self.addCleanup(patcher.stop)