JULIO-11: cross-component prototype (do not merge)

datadog_lambda/asm.py

@@ -1,4 +1,6 @@
 import logging
+import os
+import re
 import urllib.parse
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Union
@@ -9,6 +11,7 @@
 from ddtrace.internal.utils import http as http_utils
 from ddtrace.trace import Span
 
+from datadog_lambda.config import config
 from datadog_lambda.trigger import (
     EventSubtypes,
     EventTypes,
@@ -18,6 +21,32 @@
 
 logger = logging.getLogger(__name__)
 
+AWS_ACCESS_KEY_ID_TAG = "aws.lambda.access_key_id"
+_AWS_ACCESS_KEY_ID_PATTERN = re.compile(r"(ASIA|AKIA)[A-Z0-9]{16}")
+_UNSET = object()
+_aws_access_key_id: Any = _UNSET
+
+
+def _get_aws_access_key_id() -> Optional[str]:
+    """Return the execution environment's AWS access key ID, read once.
+
+    Only AWS_ACCESS_KEY_ID is read: the secret key and the session token are
+    never touched. The read happens at the first monitored invocation, so on
+    SnapStart it runs after restore, never before the snapshot.
+    """
+    global _aws_access_key_id
+    if _aws_access_key_id is _UNSET:
+        value = os.environ.get("AWS_ACCESS_KEY_ID", "")
+        if _AWS_ACCESS_KEY_ID_PATTERN.fullmatch(value):
+            _aws_access_key_id = value
+        else:
+            _aws_access_key_id = None
+            logger.debug(
+                "AWS_ACCESS_KEY_ID is empty or malformed, %s is not set",
+                AWS_ACCESS_KEY_ID_TAG,
+            )
+    return _aws_access_key_id
+
 
 def _to_single_value_headers(headers: Dict[str, List[str]]) -> Dict[str, str]:
     """
@@ -121,6 +150,11 @@ def asm_start_request(
     else:
         return
 
+    if config.appsec_aws_access_key_id_tag_enabled:
+        access_key_id = _get_aws_access_key_id()
+        if access_key_id:
+            span.set_tag(AWS_ACCESS_KEY_ID_TAG, access_key_id)
+
     body = event.get("body")
     is_base64_encoded = event.get("isBase64Encoded", False)
 

datadog_lambda/config.py

@@ -102,6 +102,9 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False):
     dsm_exchange_name = _get_env("DD_DSM_EXCHANGE_NAME")
     appsec_enabled = _get_env("DD_APPSEC_ENABLED", "false", as_bool)
     sca_enabled = _get_env("DD_APPSEC_SCA_ENABLED", "false", as_bool)
+    appsec_aws_access_key_id_tag_enabled = _get_env(
+        "DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "false", as_bool
+    )
 
     is_gov_region = _get_env("AWS_REGION", "", lambda x: x.startswith("us-gov-"))
 

tests/test_asm.py

@@ -1,8 +1,11 @@
 import json
+import logging
+import os
 import pytest
 from unittest.mock import MagicMock, patch
 
 from datadog_lambda.asm import (
+    AWS_ACCESS_KEY_ID_TAG,
     asm_start_request,
     asm_start_response,
     get_asm_blocked_response,
@@ -427,3 +430,132 @@ def test_get_asm_blocked_response_not_blocked(
     event_source = _EventSource(event_type=EventTypes.API_GATEWAY)
     response = get_asm_blocked_response(event_source)
     assert response is None
+
+
+ACCESS_KEY_ID = "ASIAEXAMPLE000000001"
+SECRET_CANARY = "canary-secret-ee3f9c1b"
+TOKEN_CANARY = "canary-token-7d41a0c2"
+
+
+@pytest.fixture
+def aws_credentials(monkeypatch):
+    import datadog_lambda.asm as asm
+
+    monkeypatch.setattr(asm, "_aws_access_key_id", asm._UNSET)
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", ACCESS_KEY_ID)
+    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", SECRET_CANARY)
+    monkeypatch.setenv("AWS_SESSION_TOKEN", TOKEN_CANARY)
+    return monkeypatch
+
+
+def _start_request(file="lambda-url.json"):
+    with open(event_samples + file, "r") as f:
+        event = json.load(f)
+    span = MagicMock()
+    with patch("datadog_lambda.asm.core"):
+        asm_start_request(
+            span,
+            event,
+            parse_event_source(event),
+            extract_trigger_tags(event, get_mock_context()),
+        )
+    return span
+
+
+def _access_key_tag_calls(span):
+    return [
+        c for c in span.set_tag.call_args_list if c.args[0] == AWS_ACCESS_KEY_ID_TAG
+    ]
+
+
+@pytest.mark.parametrize(
+    "file",
+    [
+        "lambda-url.json",
+        "api-gateway.json",
+        "http-api.json",
+        "application-load-balancer.json",
+    ],
+)
+def test_access_key_id_tag_set_when_enabled(aws_credentials, file):
+    aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "true")
+    span = _start_request(file)
+    span.set_tag.assert_any_call(AWS_ACCESS_KEY_ID_TAG, ACCESS_KEY_ID)
+    assert len(_access_key_tag_calls(span)) == 1
+
+
+@pytest.mark.parametrize("value", [None, "false", "0"])
+def test_access_key_id_tag_off_by_default(aws_credentials, value):
+    if value is not None:
+        aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", value)
+    span = _start_request()
+    assert _access_key_tag_calls(span) == []
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "",
+        "ASIA",
+        "XXXXEXAMPLE000000001",
+        "asiaexample000000001",
+        "ASIAEXAMPLE0000000012",
+        "ASIAEXAMPLE000000001\n",
+    ],
+)
+def test_access_key_id_tag_not_set_for_malformed_value(aws_credentials, value, caplog):
+    aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "true")
+    aws_credentials.setenv("AWS_ACCESS_KEY_ID", value)
+    with caplog.at_level(logging.DEBUG, logger="datadog_lambda.asm"):
+        span = _start_request()
+    assert _access_key_tag_calls(span) == []
+    assert all(r.levelno <= logging.DEBUG for r in caplog.records)
+
+
+def test_access_key_id_tag_not_set_when_variable_missing(aws_credentials):
+    aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "true")
+    aws_credentials.delenv("AWS_ACCESS_KEY_ID")
+    span = _start_request()
+    assert _access_key_tag_calls(span) == []
+
+
+def test_access_key_id_read_once(aws_credentials):
+    aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "true")
+    first = _start_request()
+    aws_credentials.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE000000002")
+    second = _start_request()
+    first.set_tag.assert_any_call(AWS_ACCESS_KEY_ID_TAG, ACCESS_KEY_ID)
+    second.set_tag.assert_any_call(AWS_ACCESS_KEY_ID_TAG, ACCESS_KEY_ID)
+
+
+def test_access_key_id_tag_not_set_for_non_http_event(aws_credentials):
+    aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "true")
+    span = _start_request("sqs-batch.json")
+    assert _access_key_tag_calls(span) == []
+
+
+class _GuardedEnviron(dict):
+    """Fails the test if the secret access key or the session token is read"""
+
+    FORBIDDEN = ("AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN")
+
+    def __getitem__(self, key):
+        assert key not in self.FORBIDDEN, f"{key} was read"
+        return super().__getitem__(key)
+
+    def get(self, key, default=None):
+        assert key not in self.FORBIDDEN, f"{key} was read"
+        return super().get(key, default)
+
+
+def test_access_key_id_never_reads_secret_or_token(aws_credentials, caplog):
+    import datadog_lambda.asm as asm
+
+    aws_credentials.setenv("DD_APPSEC_AWS_ACCESS_KEY_ID_TAG_ENABLED", "true")
+    aws_credentials.setattr(asm.os, "environ", _GuardedEnviron(os.environ))
+    with caplog.at_level(logging.DEBUG):
+        span = _start_request()
+    span.set_tag.assert_any_call(AWS_ACCESS_KEY_ID_TAG, ACCESS_KEY_ID)
+    recorded = repr(span.mock_calls) + caplog.text
+    assert SECRET_CANARY not in recorded
+    assert TOKEN_CANARY not in recorded