lambda_function.py

import pycountry_convert as conv

def lambda_handler(event, context):
    request = event['Records'][0]['cf']['request']

    '''
    Origin-Request Trigger: Updates S3 origin region based on Viewer Country.
    Updated for Python 3.12 compatibility.
    '''

    # Mappings for Region and Domain
    continent_to_region = {
        'EU': 'eu-central-1',
        'AF': 'eu-central-1',
        'NA': 'us-east-1',
        'SA': 'sa-east-1',
        'AS': 'ap-southeast-1',
        'OC': 'ap-southeast-1',
        'AN': 'ap-southeast-1'
    }

    continent_to_domain = {
        'EU': 'bahaimedia-eu.s3.eu-central-1.amazonaws.com',
        'AF': 'bahaimedia-eu.s3.eu-central-1.amazonaws.com',
        'NA': 'bahaimedia.s3.us-east-1.amazonaws.com',
        'SA': 'bahaimedia-sp.s3-sa-east-1.amazonaws.com',
        'AS': 'bahaimedia-sg.s3.ap-southeast-1.amazonaws.com',
        'OC': 'bahaimedia-sg.s3.ap-southeast-1.amazonaws.com',
        'AN': 'bahaimedia-sg.s3.ap-southeast-1.amazonaws.com'
    }

    # CloudFront headers are always lowercase
    viewer_country_header = request['headers'].get('cloudfront-viewer-country')

    if viewer_country_header:
        # CloudFront headers are a list of dicts: [{'key': '...', 'value': 'US'}]
        country_code = viewer_country_header[0]['value']

        try:
            # Attempt to convert country code (e.g., 'US') to continent (e.g., 'NA')
            continent = conv.country_alpha2_to_continent_code(country_code)

            domain_name = continent_to_domain.get(continent)
            region = continent_to_region.get(continent)

            # Fallback if continent not found in our map
            if not domain_name:
                raise ValueError("Continent not in mapping")

        except Exception as e:
            # Fallback to default US bucket on any error (conversion fail or mapping miss)
            print(f"Routing to default. Reason: {e}")
            domain_name = 'bahaimedia.s3.us-east-1.amazonaws.com'
            region = 'us-east-1'

        # Apply the routing changes to the request origin object

        # Ensure the 's3' key exists (it usually does for S3 origins, but safety first)
        if 's3' not in request['origin']:
            request['origin']['s3'] = {}

        request['origin']['s3']['region'] = region
        request['origin']['s3']['domainName'] = domain_name

        # Update the Host header so the S3 bucket accepts the request
        request['headers']['host'] = [{'key': 'host', 'value': domain_name}]

    return request