# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

#!/usr/bin/env python3
import os

import aws_cdk as cdk
import cdk_nag

from aws_cdk import Aspects
from cdk.pjwt_with_pkce_stack import PjwtWithPkceStack
from cdk.pjwt_stack import PjwtStack


app = cdk.App()

Aspects.of(app).add(cdk_nag.AwsSolutionsChecks())

stack_name = app.node.try_get_context("stack_name")

if app.node.try_get_context("pkce"):
    PjwtWithPkceStack(app, stack_name)
else:
    PjwtStack(app, stack_name)


app.synth()

app.py
aws-samples/cognito-external-idp-proxy
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import cdk_nag

from aws_cdk import (
    Aws as _aws,
    CfnOutput as _output,
    Duration,
    Stack,
    aws_apigatewayv2 as _apigw,
    aws_cognito as _cognito,
    aws_iam as _iam,
    aws_lambda as _lambda,
    aws_logs as _logs,
    aws_secretsmanager as _secretsmanager
)

from constructs import Construct
from cdk_nag import NagSuppressions

class PjwtStack(Stack):

    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)


        """
        GLOBAL CONFIGURATION ITEMS
        """

        api_version = self.node.try_get_context("api_version")
        authn_route = self.node.try_get_context("api_authn_route")
        callb_route = self.node.try_get_context("api_callback_route")
        token_route = self.node.try_get_context("api_token_route")


        """
        RESOURCE DEFINITIONS
        """

        # API Gateway base configuration - defining here to use API ID in building env vars for Lambda functions and Cognito Domain
        apigw_proxy_api = _apigw.CfnApi(
            self, "ApiGwProxyApi",
            description = "Handles requests and responses between Cognito and 3rd party IdP",
            name = construct_id + "Api",
            protocol_type = "HTTP",
            cors_configuration = _apigw.CfnApi.CorsProperty(
                allow_methods = ["POST"],
                allow_origins = ["*"]
            )
        )

        # Secrets Manager empty secret to hold the private key for private key JWT token request
        secretsmanager_private_key = _secretsmanager.Secret(
            self, "PrivateKey"
        )

        # Lambda Execution role for the token function
        lambda_function_token_exec_role = _iam.Role(
            self, "TokenRole",
            assumed_by = _iam.ServicePrincipal("lambda.amazonaws.com"),
            managed_policies = [
                _iam.ManagedPolicy.from_aws_managed_policy_name("service-role/AWSLambdaBasicExecutionRole")
            ],
            inline_policies = {
                "SecretsManagerPolicy": _iam.PolicyDocument(
                    statements = [
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["secretsmanager:GetSecretValue"],
                            resources = [secretsmanager_private_key.secret_full_arn]
                        )
                    ]
                )
            }
        )

        # Lambda function to handle token requests
        lambda_function_token = _lambda.Function(
            self, "Token",
            runtime=_lambda.Runtime.PYTHON_3_10,
            code=_lambda.Code.from_asset("./lambda/token"),
            handler="token_flow.handler",
            timeout=Duration.seconds(30),
            environment = {
                "ClientId": self.node.try_get_context("idp_client_id"),
                "ClientSecret": self.node.try_get_context("idp_client_secret"),
                "IdpIssuerUrl": self.node.try_get_context("idp_issuer_url"),
                "IdpTokenPath": self.node.try_get_context("idp_token_path"),
                "ResponseUri": f"https://{apigw_proxy_api.attr_api_id}.auth.{_aws.REGION}.amazoncognito.com/oauth2/idpresponse",
                "Pkce": str(self.node.try_get_context("pkce")),
                "Region": _aws.REGION,
                "SecretsManagerPrivateKey": secretsmanager_private_key.secret_name
            },
            role = lambda_function_token_exec_role,
            layers = [
                _lambda.LayerVersion(
                    self, "JwtLayer",
                    code = _lambda.Code.from_asset("./layers/token/"),
                    compatible_runtimes = [_lambda.Runtime.PYTHON_3_10]
                )
            ],
            log_retention = _logs.RetentionDays.FIVE_DAYS
        )

        # API Gateway Lambda integration for token
        apigw_proxy_integration_token= _apigw.CfnIntegration(
            self, "ApiGwProxyIntegrationToken",
            api_id = apigw_proxy_api.attr_api_id,
            integration_type = "AWS_PROXY",
            integration_uri = lambda_function_token.function_arn,
            integration_method = "POST",
            payload_format_version = "2.0"
        )

        # Cloudwatch log group for API Gateway logs
        cloudwatch_apigw_proxy_log_group = _logs.LogGroup(
            self, "ApiGwLogs"
        )

        # API Gateway stage
        apigw_proxy_stage = _apigw.CfnStage(
            self, "ApiGwProxyStage",
            api_id = apigw_proxy_api.attr_api_id,
            stage_name = api_version,
            access_log_settings = _apigw.CfnStage.AccessLogSettingsProperty(
                destination_arn = cloudwatch_apigw_proxy_log_group.log_group_arn,
                format = (
                    '{'
                        '"requestId": "$context.requestId",'
                        '"path": "$context.path",'
                        '"routeKey": "$context.routeKey",'
                        '"ip": "$context.identity.sourceIp",'
                        '"requestTime": "$context.requestTime",'
                        '"httpMethod": "$context.httpMethod",'
                        '"statusCode": "$context.status"'
                    '}'
                )
            ),
            auto_deploy = True
        )

        # API Gateway route for token requests between Cognito and the IdP
        apigw_proxy_route_token = _apigw.CfnRoute(
            self, "ApiGwProxyRouteToken",
            api_id = apigw_proxy_api.attr_api_id,
            route_key = f"POST {token_route}",
            target = "integrations/" + apigw_proxy_integration_token.ref
        )
        apigw_proxy_route_token_uri = apigw_proxy_api.attr_api_endpoint + "/" + apigw_proxy_stage.stage_name + token_route

        # API Gateway token route Lambda invoke permission
        lambda_function_token.add_permission(
            "ApiGwTokenRoutePermission",
            principal = _iam.ServicePrincipal("apigateway.amazonaws.com"),
            source_arn = f"arn:aws:execute-api:{_aws.REGION}:{_aws.ACCOUNT_ID}:{apigw_proxy_api.ref}/*/*{token_route}"
        )

        # API Gateway deployment definition
        apigw_proxy_deployment = _apigw.CfnDeployment(
            self, "ApiGwProxyDeployment",
            api_id = apigw_proxy_api.attr_api_id
        )
        apigw_proxy_deployment.add_dependency(apigw_proxy_route_token)

        # Translating list of scopes to cdk attributes
        cognito_oauth_scopes = []
        for scope in self.node.try_get_context("idp_scopes").split():
            if scope.lower() == "openid":
                cognito_oauth_scopes.append(_cognito.OAuthScope.OPENID)
            if scope.lower() == "email":
                cognito_oauth_scopes.append(_cognito.OAuthScope.EMAIL)
            if scope.lower() == "profile":
                cognito_oauth_scopes.append(_cognito.OAuthScope.PROFILE)
            if scope.lower() == "phone":
                cognito_oauth_scopes.append(_cognito.OAuthScope.PHONE)

        # Cognito User Pool base configuration
        cognito_user_pool = _cognito.UserPool(self, "UserPool")

        # Cognito external identity provider
        idp_issuer_url = self.node.try_get_context("idp_issuer_url")
        cognito_user_pool_idp_oidc = _cognito.UserPoolIdentityProviderOidc(
            self, "IdentityProvider",
            client_id = self.node.try_get_context("idp_client_id"),
            client_secret = self.node.try_get_context("idp_client_secret"),
            issuer_url = idp_issuer_url,
            user_pool = cognito_user_pool,
            attribute_request_method = _cognito.OidcAttributeRequestMethod.GET,
            endpoints = _cognito.OidcEndpoints(
                authorization = idp_issuer_url +self.node.try_get_context("idp_auth_path"),
                jwks_uri = idp_issuer_url + self.node.try_get_context("idp_keys_path"),
                token = apigw_proxy_route_token_uri,
                user_info = idp_issuer_url + self.node.try_get_context("idp_attributes_path")
            ),
            name = self.node.try_get_context("idp_name"),
            scopes = self.node.try_get_context("idp_scopes").split()
        )

        # Cognito Application Client
        cognito_user_pool_client = cognito_user_pool.add_client(
            "AppClient",
            o_auth = _cognito.OAuthSettings(
                flows = _cognito.OAuthFlows(
                    authorization_code_grant = True
                ),
                scopes = cognito_oauth_scopes,
                callback_urls = [self.node.try_get_context("userpool_allowed_callback_url")]
            ),
            supported_identity_providers = [_cognito.UserPoolClientIdentityProvider.custom(cognito_user_pool_idp_oidc.provider_name)]
        )

        # Cognito User Pool Domain for Hosted UI
        cognito_user_pool_domain = cognito_user_pool.add_domain(
            "CognitoDomain",
            cognito_domain = _cognito.CognitoDomainOptions(
                domain_prefix = apigw_proxy_api.attr_api_id
            )
        )


        """
        STACK OUTPUTS
        """

        _output(
            self, "ApiGwTokenEndpoint",
            value = apigw_proxy_route_token_uri
        )

        _output(
            self, "SecretsManagerPrivateKeyArn",
            value = secretsmanager_private_key.secret_full_arn
        )

        _output(
            self, "SecretsManagerPrivateKeyName",
            value = secretsmanager_private_key.secret_name
        )

        _output(
            self, "CognitoHostedUi",
            value = cognito_user_pool_domain.base_url()
        )

        _output(
            self, "CognitoIdpResponseUri",
            value = cognito_user_pool_domain.base_url() + "/oauth2/idpresponse"
        )


        """
        SUPPRESSION RULES FOR CDK_NAG
        """

        NagSuppressions.add_resource_suppressions(
            secretsmanager_private_key, [
                { "id": "AwsSolutions-SMG4", "reason": "Cannot rotate due to 3rd party IdP dependency."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            lambda_function_token_exec_role, [
                { "id": "AwsSolutions-IAM4", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions_by_path(
            self,
            f"/{self.stack_name}/LogRetentionaae0aa3c5b4d4f87b02d85b201efdd8a/ServiceRole/Resource",
            [
                {
                    "id": "AwsSolutions-IAM4",
                    "reason": "Construct specific: Enabling log retention creates a separate Lambda Function with managed policy."
                }
            ]
        )

        NagSuppressions.add_resource_suppressions_by_path(
            self,
            f"/{self.stack_name}/LogRetentionaae0aa3c5b4d4f87b02d85b201efdd8a/ServiceRole/DefaultPolicy/Resource",
            [
                {
                    "id": "AwsSolutions-IAM5",
                    "reason": "Construct specific: Enabling log retention creates a separate Lambda Function with managed policy."
                }
            ]
        )

        NagSuppressions.add_resource_suppressions(
            lambda_function_token, [
                { "id": "AwsSolutions-L1", "reason": "No tests in place to guarantee code runs in other versions."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            apigw_proxy_route_token, [
                { "id": "AwsSolutions-APIG4", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            cognito_user_pool, [
                { "id": "AwsSolutions-COG1", "reason": "Demo is supposed to integrate only with external IdP."},
                { "id": "AwsSolutions-COG2", "reason": "Defined by external IdP."},
                { "id": "AwsSolutions-COG3", "reason": "Demo purposes only."}
            ]
        )

cdk/pjwt_stack.py
aws-samples/cognito-external-idp-proxy
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import cdk_nag

from aws_cdk import (
    Aws as _aws,
    CfnOutput as _output,
    Duration,
    Stack,
    aws_apigatewayv2 as _apigw,
    aws_cognito as _cognito,
    aws_dynamodb as _dynamodb,
    aws_iam as _iam,
    aws_lambda as _lambda,
    aws_logs as _logs,
    aws_secretsmanager as _secretsmanager
)

from constructs import Construct
from cdk_nag import NagSuppressions

class PjwtWithPkceStack(Stack):

    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)


        """
        GLOBAL CONFIGURATION ITEMS
        """

        api_version = self.node.try_get_context("api_version")
        authn_route = self.node.try_get_context("api_authn_route")
        callb_route = self.node.try_get_context("api_callback_route")
        token_route = self.node.try_get_context("api_token_route")


        """
        RESOURCE DEFINITIONS
        """

        # API Gateway base configuration - defining here to use API ID in building env vars for Lambda functions and Cognito Domain
        apigw_proxy_api = _apigw.CfnApi(
            self, "ApiGwProxyApi",
            description = "Handles requests and responses between Cognito and 3rd party IdP",
            name = construct_id + "Api",
            protocol_type = "HTTP",
            cors_configuration = _apigw.CfnApi.CorsProperty(
                allow_methods = ["GET", "POST"],
                allow_origins = ["*"]
            )
        )

        # DynamoDB table to store state id and generated code verifier during authorization (PKCE)
        dynamodb_state_table = _dynamodb.Table(
            self, "StateTable",
            partition_key = _dynamodb.Attribute(name = "state", type = _dynamodb.AttributeType.STRING),
            billing_mode = _dynamodb.BillingMode.PAY_PER_REQUEST,
            time_to_live_attribute = "ttl"
        )

        # DynamoDB table to store authorization code with the associated code verfier during IdP authN response (PKCE)
        dynamodb_code_table = _dynamodb.Table(
            self, "CodeTable",
            partition_key = _dynamodb.Attribute(name = "auth_code", type = _dynamodb.AttributeType.STRING),
            billing_mode = _dynamodb.BillingMode.PAY_PER_REQUEST,
            time_to_live_attribute = "ttl"
        )

        # Secrets Manager empty secret to hold the private key for private key JWT token request
        secretsmanager_private_key = _secretsmanager.Secret(
            self, "PrivateKey"
        )

        # Lambda Execution role for the authorize function
        lambda_function_authorize_exec_role = _iam.Role(
            self, "AuthorizeRole",
            assumed_by = _iam.ServicePrincipal("lambda.amazonaws.com"),
            managed_policies = [
                _iam.ManagedPolicy.from_aws_managed_policy_name("service-role/AWSLambdaBasicExecutionRole")
            ],
            inline_policies = {
                "DynamoDbPolicy": _iam.PolicyDocument(
                    statements = [
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["dynamodb:DescribeTable", "dynamodb:PutItem"],
                            resources = [dynamodb_state_table.table_arn]
                        )
                    ]
                ),
                "SecretsManagerPolicy": _iam.PolicyDocument(
                    statements = [
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["secretsmanager:GetRandomPassword"],
                            resources = ["*"]
                        )
                    ]
                )
            }
        )

        # Lambda function to handle authorization requests when using PKCE
        lambda_function_authorize = _lambda.Function(
            self, "Authorize",
            runtime = _lambda.Runtime.PYTHON_3_10,
            code = _lambda.Code.from_asset("./lambda/authorize"),
            handler = "authorize_flow.handler",
            environment = {
                "ClientId": self.node.try_get_context("idp_client_id"),
                "DynamoDbStateTable": dynamodb_state_table.table_name,
                "IdpAuthUri": self.node.try_get_context("idp_issuer_url") + self.node.try_get_context("idp_auth_path"),
                "ProxyCallbackUri": f"{apigw_proxy_api.attr_api_endpoint}/{api_version}{callb_route}"
            },
            role = lambda_function_authorize_exec_role,
            log_retention = _logs.RetentionDays.FIVE_DAYS
        )

        # Lambda Execution role for the callback function
        lambda_function_callback_exec_role = _iam.Role(
            self, "CallbackRole",
            assumed_by = _iam.ServicePrincipal("lambda.amazonaws.com"),
            managed_policies = [
                _iam.ManagedPolicy.from_aws_managed_policy_name("service-role/AWSLambdaBasicExecutionRole")
            ],
            inline_policies = {
                "DynamoDbPolicy": _iam.PolicyDocument(
                    statements = [
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["dynamodb:DescribeTable", "dynamodb:GetItem"],
                            resources = [dynamodb_state_table.table_arn]
                        ),
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["dynamodb:DescribeTable", "dynamodb:PutItem"],
                            resources = [dynamodb_code_table.table_arn]
                        )
                    ]
                )
            }
        )

        # Lambda function to handle the IdP's callback when using PKCE
        lambda_function_callback = _lambda.Function(
            self, "Callback",
            runtime=_lambda.Runtime.PYTHON_3_10,
            code=_lambda.Code.from_asset("./lambda/callback"),
            handler="callback_flow.handler",
            environment = {
                "CognitoIdpResponseUri": f"https://{apigw_proxy_api.attr_api_id}.auth.{_aws.REGION}.amazoncognito.com/oauth2/idpresponse",
                "DynamoDbCodeTable": dynamodb_code_table.table_name,
                "DynamoDbStateTable": dynamodb_state_table.table_name
            },
            role = lambda_function_callback_exec_role,
            log_retention = _logs.RetentionDays.FIVE_DAYS
        )

        # Lambda Execution role for the token function
        lambda_function_token_exec_role = _iam.Role(
            self, "TokenRole",
            assumed_by = _iam.ServicePrincipal("lambda.amazonaws.com"),
            managed_policies = [
                _iam.ManagedPolicy.from_aws_managed_policy_name("service-role/AWSLambdaBasicExecutionRole")
            ],
            inline_policies = {
                "DynamoDbPolicy": _iam.PolicyDocument(
                    statements = [
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["dynamodb:DescribeTable", "dynamodb:GetItem"],
                            resources = [dynamodb_code_table.table_arn]
                        )
                    ]
                ),
                "SecretsManagerPolicy": _iam.PolicyDocument(
                    statements = [
                        _iam.PolicyStatement(
                            effect = _iam.Effect.ALLOW,
                            actions = ["secretsmanager:GetSecretValue"],
                            resources = [secretsmanager_private_key.secret_full_arn]
                        )
                    ]
                )
            }
        )

        # Lambda function to handle token requests
        lambda_function_token = _lambda.Function(
            self, "Token",
            runtime=_lambda.Runtime.PYTHON_3_10,
            code=_lambda.Code.from_asset("./lambda/token"),
            handler="token_flow.handler",
            timeout=Duration.seconds(30),
            environment = {
                "ClientId": self.node.try_get_context("idp_client_id"),
                "ClientSecret": self.node.try_get_context("idp_client_secret"),
                "DynamoDbCodeTable": dynamodb_code_table.table_name,
                "IdpIssuerUrl": self.node.try_get_context("idp_issuer_url"),
                "IdpTokenPath": self.node.try_get_context("idp_token_path"),
                "ResponseUri": f"{apigw_proxy_api.attr_api_endpoint}/{api_version}{callb_route}",
                "Pkce": str(self.node.try_get_context("pkce")),
                "Region": _aws.REGION,
                "SecretsManagerPrivateKey": secretsmanager_private_key.secret_name
            },
            role = lambda_function_token_exec_role,
            layers = [
                _lambda.LayerVersion(
                    self, "JwtLayer",
                    code = _lambda.Code.from_asset("./layers/token/"),
                    compatible_runtimes = [_lambda.Runtime.PYTHON_3_10]
                )
            ],
            log_retention = _logs.RetentionDays.FIVE_DAYS
        )

        # API Gateway Lambda integration for authorization
        apigw_proxy_integration_auth = _apigw.CfnIntegration(
            self, "ApiGwProxyIntegrationAuth",
            api_id = apigw_proxy_api.attr_api_id,
            integration_type = "AWS_PROXY",
            integration_uri = lambda_function_authorize.function_arn,
            integration_method = "GET",
            payload_format_version = "2.0"
        )

        # API Gateway Lambda integration for callback
        apigw_proxy_integration_callback= _apigw.CfnIntegration(
            self, "ApiGwProxyIntegrationCallback",
            api_id = apigw_proxy_api.attr_api_id,
            integration_type = "AWS_PROXY",
            integration_uri = lambda_function_callback.function_arn,
            integration_method = "GET",
            payload_format_version = "2.0"
        )

        # API Gateway Lambda integration for token
        apigw_proxy_integration_token= _apigw.CfnIntegration(
            self, "ApiGwProxyIntegrationToken",
            api_id = apigw_proxy_api.attr_api_id,
            integration_type = "AWS_PROXY",
            integration_uri = lambda_function_token.function_arn,
            integration_method = "POST",
            payload_format_version = "2.0"
        )

        # Cloudwatch log group for API Gateway logs
        cloudwatch_apigw_proxy_log_group = _logs.LogGroup(
            self, "ApiGwLogs"
        )

        # API Gateway stage
        apigw_proxy_stage = _apigw.CfnStage(
            self, "ApiGwProxyStage",
            api_id = apigw_proxy_api.attr_api_id,
            stage_name = api_version,
            access_log_settings = _apigw.CfnStage.AccessLogSettingsProperty(
                destination_arn = cloudwatch_apigw_proxy_log_group.log_group_arn,
                format = (
                    '{'
                        '"requestId": "$context.requestId",'
                        '"path": "$context.path",'
                        '"routeKey": "$context.routeKey",'
                        '"ip": "$context.identity.sourceIp",'
                        '"requestTime": "$context.requestTime",'
                        '"httpMethod": "$context.httpMethod",'
                        '"statusCode": "$context.status"'
                    '}'
                )
            ),
            auto_deploy = True
        )

        # API Gateway route for authorization
        apigw_proxy_route_auth = _apigw.CfnRoute(
            self, "ApiGwProxyRouteAuth",
            api_id = apigw_proxy_api.attr_api_id,
            route_key = f"GET {authn_route}",
            target = "integrations/" + apigw_proxy_integration_auth.ref
        )
        apigw_proxy_route_auth_uri = apigw_proxy_api.attr_api_endpoint + "/" + apigw_proxy_stage.stage_name + authn_route

        # API Gateway authorization route Lambda invoke permission
        lambda_function_authorize.add_permission(
            "ApiGwAuthorizeRoutePermission",
            principal = _iam.ServicePrincipal("apigateway.amazonaws.com"),
            source_arn = f"arn:aws:execute-api:{_aws.REGION}:{_aws.ACCOUNT_ID}:{apigw_proxy_api.ref}/*/*{authn_route}"
        )

        # API Gateway route for handling the IdP's callback
        apigw_proxy_route_callback = _apigw.CfnRoute(
            self, "ApiGwProxyRouteCallback",
            api_id = apigw_proxy_api.attr_api_id,
            route_key = f"GET {callb_route}",
            target = "integrations/" + apigw_proxy_integration_callback.ref
        )
        apigw_proxy_route_callb_uri = apigw_proxy_api.attr_api_endpoint + "/" + apigw_proxy_stage.stage_name + callb_route

        # API Gateway callback route Lambda invoke permission
        lambda_function_callback.add_permission(
            "ApiGwCallbackRoutePermission",
            principal = _iam.ServicePrincipal("apigateway.amazonaws.com"),
            source_arn = f"arn:aws:execute-api:{_aws.REGION}:{_aws.ACCOUNT_ID}:{apigw_proxy_api.ref}/*/*{callb_route}"
        )

        # API Gateway route for token requests between Cognito and the IdP
        apigw_proxy_route_token = _apigw.CfnRoute(
            self, "ApiGwProxyRouteToken",
            api_id = apigw_proxy_api.attr_api_id,
            route_key = f"POST {token_route}",
            target = "integrations/" + apigw_proxy_integration_token.ref
        )
        apigw_proxy_route_token_uri = apigw_proxy_api.attr_api_endpoint + "/" + apigw_proxy_stage.stage_name + token_route

        # API Gateway callback route Lambda invoke permission
        lambda_function_token.add_permission(
            "ApiGwTokenRoutePermission",
            principal = _iam.ServicePrincipal("apigateway.amazonaws.com"),
            source_arn = f"arn:aws:execute-api:{_aws.REGION}:{_aws.ACCOUNT_ID}:{apigw_proxy_api.ref}/*/*{token_route}"
        )

        # API Gateway deployment definition
        apigw_proxy_deployment = _apigw.CfnDeployment(
            self, "ApiGwProxyDeployment",
            api_id = apigw_proxy_api.attr_api_id
        )
        apigw_proxy_deployment.add_dependency(apigw_proxy_route_auth)
        apigw_proxy_deployment.add_dependency(apigw_proxy_route_callback)
        apigw_proxy_deployment.add_dependency(apigw_proxy_route_token)

        # Translating list of scopes to cdk attributes
        cognito_oauth_scopes = []
        for scope in self.node.try_get_context("idp_scopes").split():
            if scope.lower() == "openid":
                cognito_oauth_scopes.append(_cognito.OAuthScope.OPENID)
            if scope.lower() == "email":
                cognito_oauth_scopes.append(_cognito.OAuthScope.EMAIL)
            if scope.lower() == "profile":
                cognito_oauth_scopes.append(_cognito.OAuthScope.PROFILE)
            if scope.lower() == "phone":
                cognito_oauth_scopes.append(_cognito.OAuthScope.PHONE)

        # Cognito User Pool base configuration
        cognito_user_pool = _cognito.UserPool(self, "UserPool")

        # Cognito external identity provider
        idp_issuer_url = self.node.try_get_context("idp_issuer_url")
        cognito_user_pool_idp_oidc = _cognito.UserPoolIdentityProviderOidc(
            self, "IdentityProvider",
            client_id = self.node.try_get_context("idp_client_id"),
            client_secret = self.node.try_get_context("idp_client_secret"),
            issuer_url = idp_issuer_url,
            user_pool = cognito_user_pool,
            attribute_request_method = _cognito.OidcAttributeRequestMethod.GET,
            endpoints = _cognito.OidcEndpoints(
                authorization = apigw_proxy_route_auth_uri,
                jwks_uri = idp_issuer_url + self.node.try_get_context("idp_keys_path"),
                token = apigw_proxy_route_token_uri,
                user_info = idp_issuer_url + self.node.try_get_context("idp_attributes_path")
            ),
            name = self.node.try_get_context("idp_name"),
            scopes = self.node.try_get_context("idp_scopes").split()
        )

        # Cognito Application Client
        cognito_user_pool_client = cognito_user_pool.add_client(
            "AppClient",
            o_auth = _cognito.OAuthSettings(
                flows = _cognito.OAuthFlows(
                    authorization_code_grant = True
                ),
                scopes = cognito_oauth_scopes,
                callback_urls = [self.node.try_get_context("userpool_allowed_callback_url")]
            ),
            supported_identity_providers = [_cognito.UserPoolClientIdentityProvider.custom(cognito_user_pool_idp_oidc.provider_name)]
        )

        # Cognito User Pool Domain for Hosted UI
        cognito_user_pool_domain = cognito_user_pool.add_domain(
            "CognitoDomain",
            cognito_domain = _cognito.CognitoDomainOptions(
                domain_prefix = apigw_proxy_api.attr_api_id
            )
        )


        """
        STACK OUTPUTS
        """

        _output(
            self, "ApiGwAuthnEndpoint",
            value = apigw_proxy_route_auth_uri
        )

        _output(
            self, "ApiGwCallbackEndpoint",
            value = apigw_proxy_route_callb_uri
        )

        _output(
            self, "ApiGwTokenEndpoint",
            value = apigw_proxy_route_token_uri
        )

        _output(
            self, "DynamoDbStateTable",
            value = dynamodb_state_table.table_name
        )

        _output(
            self, "DynamoDbCodeTable",
            value = dynamodb_code_table.table_name
        )

        _output(
            self, "SecretsManagerPrivateKeyArn",
            value = secretsmanager_private_key.secret_full_arn
        )

        _output(
            self, "SecretsManagerPrivateKeyName",
            value = secretsmanager_private_key.secret_name
        )

        _output(
            self, "CognitoHostedUi",
            value = cognito_user_pool_domain.base_url()
        )

        _output(
            self, "CognitoIdpResponseUri",
            value = cognito_user_pool_domain.base_url() + "/oauth2/idpresponse"
        )


        """
        SUPPRESSION RULES FOR CDK_NAG
        """

        NagSuppressions.add_resource_suppressions(
            secretsmanager_private_key, [
                { "id": "AwsSolutions-SMG4", "reason": "Cannot rotate due to 3rd party IdP dependency."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            lambda_function_authorize_exec_role, [
                { "id": "AwsSolutions-IAM4", "reason": "Demo purposes only." },
                { "id": "AwsSolutions-IAM5", "reason": "API is resource agnostic but resource key required in statement." },
            ]
        )


        NagSuppressions.add_resource_suppressions(
            lambda_function_callback_exec_role, [
                { "id": "AwsSolutions-IAM4", "reason": "Demo purposes only."}
            ]
        )


        NagSuppressions.add_resource_suppressions(
            lambda_function_token_exec_role, [
                { "id": "AwsSolutions-IAM4", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions_by_path(
            self,
            f"/{self.stack_name}/LogRetentionaae0aa3c5b4d4f87b02d85b201efdd8a/ServiceRole/Resource",
            [
                {
                    "id": "AwsSolutions-IAM4",
                    "reason": "Construct specific: Enabling log retention creates a separate Lambda Function with managed policy."
                }
            ]
        )

        NagSuppressions.add_resource_suppressions_by_path(
            self,
            f"/{self.stack_name}/LogRetentionaae0aa3c5b4d4f87b02d85b201efdd8a/ServiceRole/DefaultPolicy/Resource",
            [
                {
                    "id": "AwsSolutions-IAM5",
                    "reason": "Construct specific: Enabling log retention creates a separate Lambda Function with managed policy."
                }
            ]
        )

        NagSuppressions.add_resource_suppressions(
            lambda_function_authorize, [
                { "id": "AwsSolutions-L1", "reason": "No tests in place to guarantee code runs in other versions."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            lambda_function_callback, [
                { "id": "AwsSolutions-L1", "reason": "No tests in place to guarantee code runs in other versions."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            lambda_function_token, [
                { "id": "AwsSolutions-L1", "reason": "No tests in place to guarantee code runs in other versions."}
            ]
        )


        NagSuppressions.add_resource_suppressions(
            apigw_proxy_route_auth, [
                { "id": "AwsSolutions-APIG4", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            apigw_proxy_route_callback, [
                { "id": "AwsSolutions-APIG4", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            apigw_proxy_route_token, [
                { "id": "AwsSolutions-APIG4", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            cognito_user_pool, [
                { "id": "AwsSolutions-COG1", "reason": "Demo is supposed to integrate only with external IdP."},
                { "id": "AwsSolutions-COG2", "reason": "Defined by external IdP."},
                { "id": "AwsSolutions-COG3", "reason": "Demo purposes only."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            dynamodb_state_table, [
                { "id": "AwsSolutions-DDB3", "reason": "Short lived data only."}
            ]
        )

        NagSuppressions.add_resource_suppressions(
            dynamodb_code_table, [
                { "id": "AwsSolutions-DDB3", "reason": "Short lived data only."}
            ]
        )

cdk/pjwt_with_pkce_stack.py
aws-samples/cognito-external-idp-proxy
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import base64
import boto3
import hashlib
import json
import os
import re
import secrets
import string
import time
import urllib.parse


"""
Initializing SDK clients to facilitate reuse across executions
"""

SM_CLIENT = boto3.client('secretsmanager')
DYNAMODB_CLIENT = boto3.client("dynamodb")


def validate_request(params: dict) -> bool:
    """
    Helper function to validate request parameters - can be used to drop requests early during runtime.
    """
    validation = False

    if params["client_id"] == os.environ.get("ClientId"):
        validation = True

    return validation


def handler(event, context):
    print("+++ FULL EVENT DETAILS +++")
    print(event)

    # Collecting original request details
    params = event["queryStringParameters"]
    org_headers = event["headers"]

    # All print statements are left in to observe behaviour in CloudWatch
    print("+++ ORIGINAL REQUEST HEADER +++")
    print(org_headers)
    print("+++ ORIGINAL QUERY STRING PARAMs")
    print(event["queryStringParameters"])
    print("#####################")

    if not validate_request(params):
        print("+++ VALIDATION FAILED - CANCELLING +++")
        return { "statusCode": 400 }

    print("+++ VALIDATION SUCCESSFUL - PROCEEDING +++")

    # Collecting envvars and necessary original request parameters
    config = {}
    config["idp_auth_uri"] = os.environ.get("IdpAuthUri")
    config["proxy_callback_uri"] = os.environ.get("ProxyCallbackUri")
    config["state_table"] = os.environ.get("DynamoDbStateTable")

    print("+++ ENV VARS COLLECTED +++")
    print(config)

    # Generate code_verifier, hash it and remove padding
    code_verifier = SM_CLIENT.get_random_password(
        PasswordLength=64,
        ExcludePunctuation=True,
        IncludeSpace=False
    )
    code_verifier = code_verifier["RandomPassword"]

    code_challenge = hashlib.sha256(code_verifier.encode("utf-8")).digest()
    code_challenge = base64.urlsafe_b64encode(code_challenge).decode('utf-8')
    code_challenge = code_challenge.replace('=', '')

    # store hashed state and code_verifier in dynamodb with ttl
    hashed_state = hashlib.sha256(params["state"].encode("utf-8")).hexdigest()
    print(f"Storing hashed state: {hashed_state}")
    state_ttl = int(time.time()) + 300
    DYNAMODB_CLIENT.put_item(
        TableName = config["state_table"],
        Item = {
            "state": {
                "S": str(hashed_state)
            },
            "code_verifier": {
                "S": code_verifier
            },
            "ttl": {
                "N": str(state_ttl)
            }
        }
    )

    # add code_challenge, code_challenge_method=S256 and proxy callback uri to params
    params["code_challenge"] = code_challenge
    params["code_challenge_method"] = "S256"
    params["redirect_uri"] = config["proxy_callback_uri"]

    print("+++ NEW QUERY STRING PARAMETERS +++")
    print(params)

    # call /authorize endpoint of the IdP with new params
    idp_redirect = config["idp_auth_uri"] + "?" + urllib.parse.urlencode(params)

    print("+++ REDIRECTING TO: +++")
    print(idp_redirect)

    # 302 to the IdP's authorize endpoint
    response_to_client = {}
    response_to_client["statusCode"] = 302

    response_to_client["body"] = json.dumps(dict())
    response_to_client["headers"] = {"Location": idp_redirect}

    return response_to_client

lambda/authorize/authorize_flow.py
aws-samples/cognito-external-idp-proxy
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import boto3
import hashlib
import json
import os
import time
import urllib.parse


"""
Initializing SDK clients to facilitate reuse across executions
"""

DYNAMODB_CLIENT = boto3.client("dynamodb")


def validate_request(params: dict) -> bool:
    """
    Helper function to validate request parameters - can be used to drop requests early during runtime.
    """
    validation = False

    if params["state"] and params["code"]:
        validation = True

    return validation


def handler(event, context):
    #  All prints left in to observe behaviour in CloudWatch
    print("+++ FULL EVENT DETAILS +++")
    print(event)

    print("+++ ORIGINAL REQUEST HEADER +++")
    org_headers = event["headers"]
    print(org_headers)

    print("+++ ORIGINAL QUERY STRING PARAMs")
    org_params = event["queryStringParameters"]
    print(org_params)
    print("#####################")

    if not validate_request(org_params):
        print("+++ VALIDATION FAILED - CANCELLING +++")
        return { "statusCode": 400 }

    print("+++ VALIDATION SUCCESSFUL - PROCEEDING +++")

    # Collecting details from env vars and query string parameters
    config = {}
    config["auth_code"] = org_params["code"]
    config["state"] = org_params["state"]
    config["auth_code_table"] = os.environ.get("DynamoDbCodeTable")
    config["state_table"] = os.environ.get("DynamoDbStateTable")
    config["cognito_idp_response_uri"] = os.environ.get("CognitoIdpResponseUri")

    print("+++ CONFIGURATION ITEMS +++")
    print(config)

    # Get code_verifier from state_table with hashed state
    hashed_state = hashlib.sha256(config["state"].encode("utf-8")).hexdigest()
    print(f"Looking for state hash: {hashed_state}")
    state_result = DYNAMODB_CLIENT.get_item(
        TableName = config["state_table"],
        Key = {
            "state": {
                "S": str(hashed_state)
            }
        }
    )

    code_verifier = state_result["Item"]["code_verifier"]["S"]

    print("+++ CODE VERIFIER FOUND +++")
    print(code_verifier)

    # Store auth_code and code_verifier in auth_code_table
    code_ttl = int(time.time()) + 300
    DYNAMODB_CLIENT.put_item(
        TableName = config["auth_code_table"],
        Item = {
            "auth_code": {
                "S": config["auth_code"]
            },
            "code_verifier": {
                "S": code_verifier
            },
            "ttl": {
                "N": str(code_ttl)
            }
        }
    )

    # Switch URL to Cognito IdP response URL and attach original query string parameters
    cognito_redirect_url = config["cognito_idp_response_uri"] + "?" + urllib.parse.urlencode(org_params)
    redirect_to_cognito = {}
    redirect_to_cognito["statusCode"] = 302
    redirect_to_cognito["body"] = json.dumps(dict())
    redirect_to_cognito["headers"] = {"Location": cognito_redirect_url}

    print("+++ CRAFTING REDIRECT TO COGNITO +++")
    print(redirect_to_cognito)

    return redirect_to_cognito

lambda/callback/callback_flow.py
aws-samples/cognito-external-idp-proxy
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

from botocore.exceptions import ClientError

import base64
import boto3
import http.client
import json
import jwt
import os
import pprint
import time
import urllib.parse


"""
Initializing SDK clients to facilitate reuse across executions
"""

DYNAMODB_CLIENT = boto3.client("dynamodb")

SM_CLIENT = boto3.client(
    service_name = "secretsmanager",
    region_name = os.environ.get("AWS_REGION")
)


def validate_request(params: dict) -> bool:
    """
    Helper function to validate request parameters - can be used to drop requests early during runtime.
    """
    validation = False

    if params["client_id"] == os.environ.get("ClientId") and params["client_secret"] == os.environ.get("ClientSecret"):
        validation = True

    return validation


def get_secret(secret_name):

    # Helper function to get secret from Secrets Manager

    try:
        get_secret_value_response = SM_CLIENT.get_secret_value(
            SecretId=secret_name
        )
    except ClientError as e:
        # For a list of exceptions thrown, see
        # https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
        raise e

    secret = get_secret_value_response['SecretString']

    return secret


def handler(event, context):
    #  All prints left in to observe behaviour in CloudWatch
    print("+++ FULL EVENT DETAILS +++")
    print(event)
    print("#####################")

    # Decode the cognito request and convert to utf-8
    encoded_message = event["body"]
    decoded_message = base64.b64decode(encoded_message)
    decoded_message = decoded_message.decode("utf-8")

    print("+++ DECODED COGNITO REQUEST +++")
    print(decoded_message)

    # Create parameter dictionary from request
    param_list = list(decoded_message.split("&"))
    param_dict = {}
    for item in param_list:
        key, value = item.split("=")
        param_dict[key] = value

    print("+++ DECODED PARAMETER LIST +++")
    print(param_dict)

    if not validate_request(param_dict):
        print("+++ VALIDATION FAILED - CANCELLING +++")
        return { "statusCode": 400 }

    print("+++ VALIDATION SUCCESSFUL - PROCEEDING +++")

    # Defining pkce toggle here because it is required in multiple different parts below
    pkce_toggle = False

    if os.environ.get("Pkce").lower() == "true":
        pkce_toggle = True
        print("+++ USING PKCE +++")

    # Fetching all details from original request and env vars
    config = {}
    config["auth_code"] = param_dict["code"]
    config["client_id"] = param_dict["client_id"]
    config["idp_issuer_url"] = os.environ.get("IdpIssuerUrl")
    config["idp_token_path"] = os.environ.get("IdpTokenPath")
    config["idp_token_endpoint"] = config["idp_issuer_url"] + config["idp_token_path"]
    config["secret_name"] = os.environ.get("SecretsManagerPrivateKey")
    config["original_response_uri"] = os.environ.get("ResponseUri")

    if pkce_toggle:
        config["code_table"] = os.environ.get("DynamoDbCodeTable")

    print("+++ CONFIGURATION ITEMS +++")
    print(config)

    # Get code_verifier associated with auth_token when using PKCE
    if pkce_toggle:
        code_result = DYNAMODB_CLIENT.get_item(
            TableName = config["code_table"],
            Key = {
                "auth_code": {
                    "S": config["auth_code"]
                }
            }
        )
        code_verifier = code_result["Item"]["code_verifier"]["S"]

        print("+++ CODE VERIFIER FOUND +++")
        print(code_verifier)

    # Get private key from Secrets Manager
    print("+++ RETRIEVING SECRET FROM SECRET MANAGER +++")
    private_key = get_secret(config["secret_name"])
    private_key_dict = json.loads(private_key)
    private_key = jwt.jwk_from_dict(private_key_dict)
    print("+++ KEY RETRIEVED +++")

    print("+++ SIGNING TOKEN +++")
    # Create the private key jwt
    instance = jwt.JWT()
    private_key_jwt = instance.encode({
        "iss": config["client_id"],
        "sub": config["client_id"],
        "aud": config["idp_token_endpoint"],
        "iat": int(time.time()),
        "exp": int(time.time()) + 300
    },
        private_key,
        alg='RS256',
        optional_headers = {"kid": private_key_dict["kid"]}
    )

    print("+++ PRIVATE KEY JWT +++")
    print(private_key_jwt)

    # Add client_assertion to the query string params
    param_dict["client_assertion"] = private_key_jwt
    param_dict["grant_type"] = "authorization_code"
    param_dict["client_assertion_type"] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
    param_dict["redirect_uri"] = config["original_response_uri"]

    # Add the api gw url from the authorize request and code verifier when using PKCE
    if pkce_toggle:
        param_dict["code_verifier"] = code_verifier

    # Removing because it is not needed
    param_dict.pop("client_secret")

    # Make the token request
    clean_url = config["idp_issuer_url"][8:]
    conn = http.client.HTTPSConnection(clean_url)

    payload = urllib.parse.urlencode(param_dict)
    print("+++ PAYLOAD +++")
    print(payload)

    headers = {
        "Content-Type": "application/x-www-form-urlencoded"
    }

    conn.request('POST', f'{config["idp_token_path"]}', payload, headers)
    res = conn.getresponse()

    print("+++ IDP RESPONSE +++")
    print(f"Status: {res.status}, Reason {res.reason}")

    # Return IdP response to Cognito
    data = res.read().decode("utf-8")

    print(data)

    return data


lambda/token/token_flow.py
aws-samples/cognito-external-idp-proxy
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

from jwcrypto import jwk

import os
import uuid

# Defining output filenames for keys
pbkey_name = "public_key.json"
prkey_name = "private_key.json"
key_dir = "local"

# Create a random string as key identifier
keyid = str(uuid.uuid4())

# Generate key material for public anc private key as JWK
key = jwk.JWK.generate(kty="RSA", size=4096, kid=keyid)
public_key = key.export_public()
private_key = key.export_private()

# Store key to local filesystem and print output for quick reference
try:
    os.mkdir(key_dir)

except FileExistsError:
    print("Directory exists - using that.")

except:
    raise


try:
    with open(f"{key_dir}/{pbkey_name}", "w") as pbk:
        pbk.write(public_key)

    print("\n\n### PUBLIC KEY TO PASTE IN IDP ###\n")
    print(public_key)

    with open(f"{key_dir}/{prkey_name}", "w") as prk:
        prk.write(private_key)

    print("\n\n### PRIVATE KEY! STORE SECURELY! ###\n")
    print(private_key)

except:
    raise

finally:
    print(f"\n\nKey files generated: {key_dir}/{pbkey_name}, {key_dir}/{prkey_name}")

scripts/generate-keys.py
aws-samples/cognito-external-idp-proxy