datadog_lambda 5.91.0__py3-none-any.whl

datadog_lambda/trigger.py

@@ -0,0 +1,352 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# 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 gzip
+import json
+from io import BytesIO, BufferedReader
+from enum import Enum
+from typing import Any
+
+
+class _stringTypedEnum(Enum):
+    """
+    _stringTypedEnum provides a type-hinted convenience function for getting the string value of
+    an enum.
+    """
+
+    def get_string(self) -> str:
+        return self.value
+
+
+class EventTypes(_stringTypedEnum):
+    """
+    EventTypes is an enum of Lambda event types we care about.
+    """
+
+    UNKNOWN = "unknown"
+    API_GATEWAY = "api-gateway"
+    APPSYNC = "appsync"
+    ALB = "application-load-balancer"
+    CLOUDWATCH_LOGS = "cloudwatch-logs"
+    CLOUDWATCH_EVENTS = "cloudwatch-events"
+    CLOUDFRONT = "cloudfront"
+    DYNAMODB = "dynamodb"
+    EVENTBRIDGE = "eventbridge"
+    KINESIS = "kinesis"
+    LAMBDA_FUNCTION_URL = "lambda-function-url"
+    S3 = "s3"
+    SNS = "sns"
+    SQS = "sqs"
+    STEPFUNCTIONS = "states"
+
+
+class EventSubtypes(_stringTypedEnum):
+    """
+    EventSubtypes is an enum of Lambda event subtypes.
+    Currently, API Gateway events subtypes are supported,
+    e.g. HTTP-API and Websocket events vs vanilla API-Gateway events.
+    """
+
+    NONE = "none"
+    API_GATEWAY = "api-gateway"  # regular API Gateway
+    WEBSOCKET = "websocket"
+    HTTP_API = "http-api"
+
+
+class _EventSource:
+    """
+    _EventSource holds an event's type and subtype.
+    """
+
+    def __init__(
+        self,
+        event_type: EventTypes,
+        subtype: EventSubtypes = EventSubtypes.NONE,
+    ):
+        self.event_type = event_type
+        self.subtype = subtype
+
+    def to_string(self) -> str:
+        """
+        to_string returns the string representation of an _EventSource.
+        Since to_string was added to support trigger tagging,
+        the event's subtype will never be included in the string.
+        """
+        return self.event_type.get_string()
+
+    def equals(
+        self, event_type: EventTypes, subtype: EventSubtypes = EventSubtypes.NONE
+    ) -> bool:
+        """
+        equals provides syntactic sugar to determine whether this _EventSource has a given type
+        and subtype.
+        Unknown events will never equal other events.
+        """
+        if self.event_type == EventTypes.UNKNOWN:
+            return False
+        if self.event_type != event_type:
+            return False
+        if self.subtype != subtype:
+            return False
+        return True
+
+
+def get_aws_partition_by_region(region):
+    if region.startswith("us-gov-"):
+        return "aws-us-gov"
+    if region.startswith("cn-"):
+        return "aws-cn"
+    return "aws"
+
+
+def get_first_record(event):
+    records = event.get("Records")
+    if records and len(records) > 0:
+        return records[0]
+
+
+def parse_event_source(event: dict) -> _EventSource:
+    """Determines the source of the trigger event"""
+    if type(event) is not dict:
+        return _EventSource(EventTypes.UNKNOWN)
+
+    event_source = _EventSource(EventTypes.UNKNOWN)
+
+    request_context = event.get("requestContext")
+    if request_context and request_context.get("stage"):
+        if "domainName" in request_context and detect_lambda_function_url_domain(
+            request_context.get("domainName")
+        ):
+            return _EventSource(EventTypes.LAMBDA_FUNCTION_URL)
+        event_source = _EventSource(EventTypes.API_GATEWAY)
+        if "httpMethod" in event:
+            event_source.subtype = EventSubtypes.API_GATEWAY
+        if "routeKey" in event:
+            event_source.subtype = EventSubtypes.HTTP_API
+        if event.get("requestContext", {}).get("messageDirection"):
+            event_source.subtype = EventSubtypes.WEBSOCKET
+
+    if request_context and request_context.get("elb"):
+        event_source = _EventSource(EventTypes.ALB)
+
+    if event.get("awslogs"):
+        event_source = _EventSource(EventTypes.CLOUDWATCH_LOGS)
+
+    if event.get("detail-type"):
+        event_source = _EventSource(EventTypes.EVENTBRIDGE)
+
+    event_detail = event.get("detail")
+    has_event_categories = (
+        isinstance(event_detail, dict)
+        and event_detail.get("EventCategories") is not None
+    )
+    if event.get("source") == "aws.events" or has_event_categories:
+        event_source = _EventSource(EventTypes.CLOUDWATCH_EVENTS)
+
+    if "Execution" in event and "StateMachine" in event and "State" in event:
+        event_source = _EventSource(EventTypes.STEPFUNCTIONS)
+
+    event_record = get_first_record(event)
+    if event_record:
+        aws_event_source = event_record.get(
+            "eventSource", event_record.get("EventSource")
+        )
+
+        if aws_event_source == "aws:dynamodb":
+            event_source = _EventSource(EventTypes.DYNAMODB)
+        if aws_event_source == "aws:kinesis":
+            event_source = _EventSource(EventTypes.KINESIS)
+        if aws_event_source == "aws:s3":
+            event_source = _EventSource(EventTypes.S3)
+        if aws_event_source == "aws:sns":
+            event_source = _EventSource(EventTypes.SNS)
+        if aws_event_source == "aws:sqs":
+            event_source = _EventSource(EventTypes.SQS)
+
+        if event_record.get("cf"):
+            event_source = _EventSource(EventTypes.CLOUDFRONT)
+
+    return event_source
+
+
+def detect_lambda_function_url_domain(domain: str) -> bool:
+    #  e.g. "etsn5fibjr.lambda-url.eu-south-1.amazonaws.com"
+    domain_parts = domain.split(".")
+    if len(domain_parts) < 2:
+        return False
+    return domain_parts[1] == "lambda-url"
+
+
+def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> str:
+    """
+    Parses the trigger event for an available ARN. If an ARN field is not provided
+    in the event we stitch it together.
+    """
+    split_function_arn = context.invoked_function_arn.split(":")
+    region = split_function_arn[3]
+    account_id = split_function_arn[4]
+    aws_arn = get_aws_partition_by_region(region)
+
+    event_record = get_first_record(event)
+    # e.g. arn:aws:s3:::lambda-xyz123-abc890
+    if source.to_string() == "s3":
+        return event_record.get("s3", {}).get("bucket", {}).get("arn")
+
+    # e.g. arn:aws:sns:us-east-1:123456789012:sns-lambda
+    if source.to_string() == "sns":
+        return event_record.get("Sns", {}).get("TopicArn")
+
+    # e.g. arn:aws:cloudfront::123456789012:distribution/ABC123XYZ
+    if source.event_type == EventTypes.CLOUDFRONT:
+        distribution_id = (
+            event_record.get("cf", {}).get("config", {}).get("distributionId")
+        )
+        return "arn:{}:cloudfront::{}:distribution/{}".format(
+            aws_arn, account_id, distribution_id
+        )
+
+    # e.g. arn:aws:lambda:<region>:<account_id>:url:<function-name>:<function-qualifier>
+    if source.equals(EventTypes.LAMBDA_FUNCTION_URL):
+        function_name = ""
+        if len(split_function_arn) >= 7:
+            function_name = split_function_arn[6]
+        function_arn = f"arn:aws:lambda:{region}:{account_id}:url:{function_name}"
+        function_qualifier = ""
+        if len(split_function_arn) >= 8:
+            function_qualifier = split_function_arn[7]
+            function_arn = function_arn + f":{function_qualifier}"
+        return function_arn
+
+    # e.g. arn:aws:apigateway:us-east-1::/restapis/xyz123/stages/default
+    if source.event_type == EventTypes.API_GATEWAY:
+        request_context = event.get("requestContext")
+        return "arn:{}:apigateway:{}::/restapis/{}/stages/{}".format(
+            aws_arn, region, request_context.get("apiId"), request_context.get("stage")
+        )
+
+    # e.g. arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/lambda-xyz/123
+    if source.event_type == EventTypes.ALB:
+        request_context = event.get("requestContext")
+        return request_context.get("elb", {}).get("targetGroupArn")
+
+    # e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz
+    if source.event_type == EventTypes.CLOUDWATCH_LOGS:
+        with gzip.GzipFile(
+            fileobj=BytesIO(base64.b64decode(event.get("awslogs", {}).get("data")))
+        ) as decompress_stream:
+            data = b"".join(BufferedReader(decompress_stream))
+        logs = json.loads(data)
+        log_group = logs.get("logGroup", "cloudwatch")
+        return "arn:{}:logs:{}:{}:log-group:{}".format(
+            aws_arn, region, account_id, log_group
+        )
+
+    # e.g. arn:aws:events:us-east-1:123456789012:rule/my-schedule
+    if source.event_type == EventTypes.CLOUDWATCH_EVENTS and event.get("resources"):
+        return event.get("resources")[0]
+
+
+def get_event_source_arn(source: _EventSource, event: dict, context: Any) -> str:
+    event_source_arn = event.get("eventSourceARN") or event.get("eventSourceArn")
+
+    event_record = get_first_record(event)
+    if event_record:
+        event_source_arn = event_record.get("eventSourceARN") or event_record.get(
+            "eventSourceArn"
+        )
+
+    if event_source_arn is None:
+        event_source_arn = parse_event_source_arn(source, event, context)
+
+    return event_source_arn
+
+
+def extract_http_tags(event):
+    """
+    Extracts HTTP facet tags from the triggering event
+    """
+    http_tags = {}
+    request_context = event.get("requestContext")
+    path = event.get("path")
+    method = event.get("httpMethod")
+    if request_context and request_context.get("stage"):
+        if request_context.get("domainName"):
+            http_tags["http.url"] = request_context.get("domainName")
+
+        path = request_context.get("path")
+        method = request_context.get("httpMethod")
+        # Version 2.0 HTTP API Gateway
+        apigateway_v2_http = request_context.get("http")
+        if event.get("version") == "2.0" and apigateway_v2_http:
+            path = apigateway_v2_http.get("path")
+            method = apigateway_v2_http.get("method")
+
+    if path:
+        http_tags["http.url_details.path"] = path
+    if method:
+        http_tags["http.method"] = method
+
+    headers = event.get("headers")
+    if headers and headers.get("Referer"):
+        http_tags["http.referer"] = headers.get("Referer")
+
+    return http_tags
+
+
+def extract_trigger_tags(event: dict, context: Any) -> dict:
+    """
+    Parses the trigger event object to get tags to be added to the span metadata
+    """
+    trigger_tags = {}
+    event_source = parse_event_source(event)
+    if event_source.to_string() is not None and event_source.to_string() != "unknown":
+        trigger_tags["function_trigger.event_source"] = event_source.to_string()
+
+        event_source_arn = get_event_source_arn(event_source, event, context)
+        if event_source_arn:
+            trigger_tags["function_trigger.event_source_arn"] = event_source_arn
+
+    if event_source.event_type in [
+        EventTypes.API_GATEWAY,
+        EventTypes.ALB,
+        EventTypes.LAMBDA_FUNCTION_URL,
+    ]:
+        trigger_tags.update(extract_http_tags(event))
+
+    return trigger_tags
+
+
+def extract_http_status_code_tag(trigger_tags, response):
+    """
+    If the Lambda was triggered by API Gateway, Lambda Function URL, or ALB,
+    add the returned status code as a tag to the function execution span.
+    """
+    if trigger_tags is None:
+        return
+    str_event_source = trigger_tags.get("function_trigger.event_source")
+    # it would be cleaner if each event type was a constant object that
+    # knew some properties about itself like this.
+    str_http_triggers = [
+        et.value
+        for et in [
+            EventTypes.API_GATEWAY,
+            EventTypes.LAMBDA_FUNCTION_URL,
+            EventTypes.ALB,
+        ]
+    ]
+    if str_event_source not in str_http_triggers:
+        return
+
+    status_code = "200"
+    if response is None:
+        # Return a 502 status if no response is found
+        status_code = "502"
+    elif hasattr(response, "get"):
+        status_code = response.get("statusCode")
+    elif hasattr(response, "status_code"):
+        status_code = response.status_code
+
+    return str(status_code)