Cognito-MWAA Connector

lambda_function.py

import json
import urllib.request
import urllib.parse
import boto3

# This lambda function is configured in the AWS console to have a "function URL" 
# a dedicated https endpoint for this function. In this configuration, the 
# event parameter is a dict describing the request made to the function URL
def lambda_handler(event, context):
    
    rikolti_userpool = "https://rikolti-airflow.auth.us-east-2.amazoncognito.com"
    lambda_app_client_id = "5c2aj2bm3joeknb8667i93lbri"
    lambda_app_client_uri = "https://lktz4m5uyq3kd4fprwtozkm4xq0fihdg.lambda-url.us-east-2.on.aws/"

    # Cognito Built-In UI gives us an OAuth authorization code grant
    # if no user code grant, redirect to cognito login page
    user_code_grant = event.get('queryStringParameters', {}).get('code')
    if not user_code_grant:
        return {
            'statusCode': 302,
            'headers': {
                'Location': (
                    f"{rikolti_userpool}/login?"
                    f"client_id={lambda_app_client_id}&"
                    "response_type=code&scope=email+openid+phone&"
                    f"redirect_uri={lambda_app_client_uri}"
                )
            }
        }

    # Get oauth token from cognito user pool given a user code grant
    # https://docs.aws.amazon.com/cognito/latest/developerguide/token-endpoint.html
    url = f"{rikolti_userpool}/oauth2/token"
    data = {
        'grant_type': 'authorization_code',
        'client_id': lambda_app_client_id,
        'redirect_uri': lambda_app_client_uri,
        'code': user_code_grant
    }
    data = urllib.parse.urlencode(data)
    data = data.encode('ascii')
    token_request = urllib.request.Request(url, data=data, method='POST')
    with urllib.request.urlopen(token_request) as response:
        token = json.loads(response.read())
        # you could decode the token with something like pyjwt here
        # https://jwt.io/
        # https://medium.com/geekculture/how-to-encode-and-decode-jwt-token-using-python-f9c33de576c5
    
    
    # Get user info from cognito user pool given an oauth access token
    # https://docs.aws.amazon.com/cognito/latest/developerguide/userinfo-endpoint.html
    url = f"{rikolti_userpool}/oauth2/userInfo"
    headers = {"Authorization": f"Bearer {token.get('access_token')}"}
    user_info_request = urllib.request.Request(url, headers=headers, method="GET")
    with urllib.request.urlopen(user_info_request) as response:
        user_info = json.loads(response.read())
        username = user_info.get('username')
    
    # Get (or create) an identity id in cognito identity pool given a login pair
    # https://docs.aws.amazon.com/cognitoidentity/latest/APIReference/API_GetId.html
    identity_client = boto3.client('cognito-identity')
    pad_dsc_dev_acct = "866216109762"
    rikolti_identitypool_id = "us-east-2:675d4f25-40ae-4be5-a275-a866bb40253f"
    rikolti_userpool_id = "us-east-2_QDNIqakFj"
    identity_response = identity_client.get_id(
        AccountId=pad_dsc_dev_acct,
        IdentityPoolId=rikolti_identitypool_id,
        Logins={
            f"cognito-idp.us-east-2.amazonaws.com/{rikolti_userpool_id}": token.get('id_token')
        }
    )
    identity_id = identity_response.get('IdentityId')

    # Enhanced Flow - this works, but results in one Airflow 
    # user for the role, so every cognito user has the same Airflow user 
    # Authentication Flows - new "enhanced" flow vs. "classic flow":
    # https://docs.aws.amazon.com/cognito/latest/developerguide/authentication-flow.html

    # Get credentials from identity pool given an identity id and login pair
    # https://docs.aws.amazon.com/cognitoidentity/latest/APIReference/API_GetCredentialsForIdentity.html
    # credentials_response = identity_client.get_credentials_for_identity(
    #     IdentityId=identity_id,
    #     Logins={
    #         f"cognito-idp.us-east-2.amazonaws.com/{rikolti_userpool_id}": token.get('id_token')
    #     }
    # )
    # credentials = credentials_response.get('Credentials')
    
    # Classic Flow - instead of using cognito to get us credentials, we call sts
    # directly, with a RoleSessionName allowing for individual airflow users
    # per cognito user pool user

    # Get open id token from identity pool given an identity id and login pair
    # https://docs.aws.amazon.com/cognitoidentity/latest/APIReference/API_GetOpenIdToken.html
    open_id_response = identity_client.get_open_id_token(
        IdentityId=identity_id,
        Logins={
            f"cognito-idp.us-east-2.amazonaws.com/{rikolti_userpool_id}": token.get('id_token')
        }
    )
    open_id = open_id_response.get('Token')
    
    # Get credentials from sts given an open id token
    # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
    mwaa_role_arn = "arn:aws:iam::866216109762:role/service-role/mwaa-create-login"
    sts = boto3.client('sts')
    credentials_response = sts.assume_role_with_web_identity(
        RoleArn=mwaa_role_arn,
        RoleSessionName=username,
        WebIdentityToken=open_id,
    )
    credentials = credentials_response.get('Credentials')
    
    # https://docs.aws.amazon.com/mwaa/latest/userguide/call-mwaa-apis-web.html
    # Create a MWAA client using those temporary sts credentials
    # SecretAccessKey is just SecretKey in the get_credentials_for_identity endpoint
    mwaa = boto3.client('mwaa', 
        aws_access_key_id=credentials.get('AccessKeyId'),
        aws_secret_access_key=credentials.get('SecretAccessKey'), 
        aws_session_token=credentials.get('SessionToken')
    )

    response = mwaa.create_web_login_token(Name="RikoltiAirflowPilot")
    webServerHostName = response["WebServerHostname"]
    webToken = response["WebToken"]
    airflowUIUrl = (
        f"https://{webServerHostName}/aws_mwaa/aws-console-sso?"
        f"login=true#{webToken}"
    )

    return {
        'statusCode': 302,
        'headers': {
            'Location': airflowUIUrl
        }
    }