# 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

EVENT_SOURCES = [

"aws:dynamodb",

"aws:kinesis",

"aws:s3",

"aws:sns",

"aws:sqs",

]

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):

"""Determines the source of the trigger event

Possible Returns:

api-gateway | application-load-balancer | cloudwatch-logs |

cloudwatch-events | cloudfront | dynamodb | kinesis | s3 | sns | sqs

"""

event_source = event.get("eventSource") or event.get("EventSource")

request_context = event.get("requestContext")

if request_context and request_context.get("stage"):

event_source = "api-gateway"

if request_context and request_context.get("elb"):

event_source = "application-load-balancer"

if event.get("awslogs"):

event_source = "cloudwatch-logs"

event_detail = event.get("detail")

cw_event_categories = event_detail and event_detail.get("EventCategories")

if event.get("source") == "aws.events" or cw_event_categories:

event_source = "cloudwatch-events"

event_record = get_first_record(event)

if event_record:

event_source = event_record.get("eventSource") or event_record.get(

"EventSource"

)

if event_record.get("cf"):

event_source = "cloudfront"

if event_source in EVENT_SOURCES:

event_source = event_source.replace("aws:", "")

return event_source

def parse_event_source_arn(source, event, context):

"""

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 == "s3":

return event_record.get("s3")["bucket"]["arn"]

# e.g. arn:aws:sns:us-east-1:123456789012:sns-lambda

if source == "sns":

return event_record.get("Sns")["TopicArn"]

# e.g. arn:aws:cloudfront::123456789012:distribution/ABC123XYZ

if source == "cloudfront":

distribution_id = event_record.get("cf")["config"]["distributionId"]

return "arn:{}:cloudfront::{}:distribution/{}".format(

aws_arn, account_id, distribution_id

)

# e.g. arn:aws:apigateway:us-east-1::/restapis/xyz123/stages/default

if source == "api-gateway":

request_context = event.get("requestContext")

return "arn:{}:apigateway:{}::/restapis/{}/stages/{}".format(

aws_arn, region, request_context["apiId"], request_context["stage"]

)

# e.g. arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/lambda-xyz/123

if source == "application-load-balancer":

request_context = event.get("requestContext")

return request_context.get("elb")["targetGroupArn"]

# e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz

if source == "cloudwatch-logs":

with gzip.GzipFile(

fileobj=BytesIO(base64.b64decode(event["awslogs"]["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 == "cloudwatch-events" and event.get("resources"):

return event.get("resources")[0]

def get_event_source_arn(source, event, context):

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["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["Referer"]

return http_tags

def extract_trigger_tags(event, context):

"""

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:

trigger_tags["function_trigger.event_source"] = event_source

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 in ["api-gateway", "application-load-balancer"]:

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 or ALB add the returned status code

as a tag to the function execution span.

"""

is_http_trigger = trigger_tags and (

trigger_tags.get("function_trigger.event_source") == "api-gateway"

or trigger_tags.get("function_trigger.event_source")

== "application-load-balancer"

)

if not is_http_trigger:

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 status_code