Lambda@Edge function to that can call Lambda Function URLs with AWS_IAM

lambda_handler.py

# Based on https://medium.com/@dario_26152/restrict-access-to-lambda-functionurl-to-cloudfront-using-aws-iam-988583834705
import base64
from typing import TYPE_CHECKING

from boto3 import Session
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest

if TYPE_CHECKING:
    from typing import TypedDict
    from botocore.compat import HTTPHeaders as Boto3HTTPHeaders

    # https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/lambda-event-structure.html
    class CloudFrontHeader(TypedDict):
        key: str
        value: str

    CloudFrontRequestHeaders = dict[str, list[CloudFrontHeader]]

    class CloudFrontRequestBody(TypedDict):
        data: str

    class CloudFrontRequest(TypedDict):
        headers: CloudFrontRequestHeaders
        method: str
        uri: str
        body: CloudFrontRequestBody

    class LambdaAtEdgeCF(TypedDict):
        request: CloudFrontRequest

    class LambdaAtEdgeRecord(TypedDict):
        cf: LambdaAtEdgeCF

    class LambdaAtEdgePayload(TypedDict):
        Records: list[LambdaAtEdgeRecord]

session = Session()
credentials = session.get_credentials().get_frozen_credentials()


def cf_headers_to_dict(headers: "CloudFrontRequestHeaders") -> dict[str, str]:
    d = {}
    for k, v in headers.items():
        # This header should not be signed as it changes hop to hop
        if k == 'x-forwarded-for':
            continue
        d[v[0]['key']] = v[0]['value']
    return d


def boto3_headers_to_cf_headers(headers: "Boto3HTTPHeaders") -> "CloudFrontRequestHeaders":
    d: CloudFrontRequestHeaders = {}
    for k, v in headers.items():
        d[k.lower()] = [{'key': k, 'value': v}]
    return d


def sign_cf_request_for_lambda_function_url_with_aws_iam(cf_request: "CloudFrontRequest") -> "CloudFrontRequest":
    # Convert the CloudFront request to a format that SigV4Auth can understand
    headers = cf_headers_to_dict(cf_request['headers'])

    # kjh123kjhbrkj.lambda-url.eu-west-2.on.aws -> eu-west-2
    signer = SigV4Auth(credentials, 'lambda', region_name=headers['host'].split('.')[2])
    aws_request = AWSRequest(
        cf_request['method'],
        "https://" + headers['host'] + cf_request['uri'],
        data=base64.b64decode(cf_request['body']['data']),
        headers=headers,
    )
    # Create and inject signature
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html#add-signature-to-request
    signer.add_auth(aws_request)
    cf_request['headers'] = boto3_headers_to_cf_headers(aws_request.headers)
    return cf_request


def lambda_handler(event: "LambdaAtEdgePayload", context) -> "CloudFrontRequest":
    # https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/lambda-event-structure.html
    cf_request = event['Records'][0]['cf']['request']
    return sign_cf_request_for_lambda_function_url_with_aws_iam(
        cf_request=cf_request,
    )