Toan Tran

Posted on Sep 17, 2023

Role-based access control with AWS CDK

#rbac #cdk #aws #python

Overview

The purpose of this blog is to utilize an Amazon Cognito user pool as a user repository and enable users to authenticate and obtain a JSON Web Token (JWT) for subsequent use with API Gateway, the JWT serves the purpose of identifying the user's group affiliation, which, when mapped to an IAM policy, determines the access privileges granted to the group.

The GitHub repo is here

Prerequisite

AWS CDKv2
Python

Create a CDK Python project

To create a CDK project, we can use the following commands:

mkdir role-based-access-control
cdk init app --language python

As the project is ready to go, we will start creating some resources.

We will create a folder named resources under role_based_access_control. We will put all the AWS resources in this folder.

This is what the project structure looks like

.
├── README.md
├── app.py
├── cdk.json
├── requirements-dev.txt
├── requirements.txt
├── role_based_access_control
│   ├── __init__.py
│   ├── resources
│   └── role_based_access_control_stack.py
├── source.bat
└── tests
    ├── __init__.py
    └── unit
        ├── __init__.py
        └── test_role_based_access_control_stack.py

Create a User pool with Cognito

Under the resources folder, create cognito/congito_construct.py

from aws_cdk import aws_cognito as _cognito
from aws_cdk import aws_lambda as _lambda
from constructs import Construct


class Cognito(Construct):

    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        props: dict = {},
        **kwargs
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        env: str = props.get('env', 'dev')
        callback_domain: str = 'https://localhost.com'
        post_confirmation: _lambda.Function = props.get('post_confirmation', None)
        custom_auth: _lambda.Function = props.get('custom_auth', None)

        self.user_pool = _cognito.UserPool(
            self, f'{env}-userpool')

        if post_confirmation:
            self.user_pool.add_trigger(
                _cognito.UserPoolOperation.POST_CONFIRMATION,
                post_confirmation)

        _cognito.CfnUserPoolGroup(
            self, "admin",
            user_pool_id=self.user_pool.user_pool_id,
            group_name="admin"
        )

        _cognito.CfnUserPoolGroup(
            self, "individual",
            user_pool_id=self.user_pool.user_pool_id,
            group_name="individual"
        )

        _cognito.CfnUserPoolDomain(
            self, f"{env}-cognito-domain",
            domain='rbac-test-1',
            user_pool_id=self.user_pool.user_pool_id,
        )

        self.client = _cognito.UserPoolClient(
            self, f"{env}-cognito-client",
            user_pool=self.user_pool,
            supported_identity_providers=[
                _cognito.UserPoolClientIdentityProvider.COGNITO,
            ],
            o_auth={
                "callback_urls": [callback_domain]
            }
        )

        if custom_auth:
            custom_auth.add_environment(
                'COGNITO_APP_CLIENT_ID', self.client.user_pool_client_id)
            custom_auth.add_environment(
                'COGNITO_USER_POOL_ID', self.user_pool.user_pool_id)

The above code will create:

A Cognito user pool and client
Two user groups: admin and individual: we will create access policies for those groups later.
A lambda post-confirmation trigger: to update the user's group after their account is registered successfully.
Add enviroments varibale for lambda custom authentication.

Create Lambda

Create a subdirectory named "lambdas" within the "resources" directory, and inside the "lambdas" directory, place a file named "lambda_construct.py."

from aws_cdk import aws_lambda as _lambda
from constructs import Construct


class Lambda(Construct):

    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        props: dict = {},
        **kwargs
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        run_time: _lambda.Runtime = _lambda.Runtime.PYTHON_3_9

        self.post_confirmation = _lambda.Function(
            self, 'PostConfirmation',
            runtime=run_time,
            handler='lambda_function.lambda_handler',
            code=_lambda.Code.from_asset(
                './role_based_access_control/resources/lambdas/post_confirmation/src')
        )

        self.custom_auth = _lambda.Function(
            self, 'CustomAuth',
            runtime=run_time,
            handler='lambda_function.lambda_handler',
            code=_lambda.Code.from_asset(
                './role_based_access_control/resources/lambdas/custom_auth/src/')
        )

This construct will create 2 lambda functions that handle the post-confirmation event from Cognito and authentication for API Gateway.

Create DynamoDB

Create a subdirectory titled 'dynamodb' within the 'resources' directory. Within the 'dynamodb' directory, deposit a file named 'dynamodb_construct.py.'

from aws_cdk import aws_dynamodb as _dynamodb
from aws_cdk import aws_lambda as _lambda
from constructs import Construct


class DynamoDB(Construct):

    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        props: dict = {},
        **kwargs
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        env = props.get('env', 'dev')

        custom_auth: _lambda.Function = props.get('custom_auth', None)

        self.auth_table = _dynamodb.Table(
            self, f'{env}-auth-policy-store',
            partition_key=_dynamodb.Attribute(
                name="group",
                type=_dynamodb.AttributeType.STRING
            )
        )

        if custom_auth:
            custom_auth.add_environment(
                'TABLE_NAME', self.auth_table.table_name)
            self.auth_table.grant_read_write_data(custom_auth)

We also need to add an environment for lambda custom authentication. In this case, we put the table name.

Example policy for dynamodb record

{
 "group": "user",
 "policy": {
  "Statement": [
   {
    "Action": "execute-api:Invoke",
    "Effect": "Allow",
    "Resource": [
     "arn:aws:execute-api:*:*:*/*/GET/admin"
    ],
    "Sid": "API"
   }
  ],
  "Version": "2012-10-17"
 }
}

This policy allows the user or group it is attached to the permission to invoke HTTP GET requests on any API Gateway endpoint under the "/admin" resource path, regardless of the AWS account or stage. It is a fairly permissive policy for API Gateway access.

Create API Gateway & Integration.

Create a 'api_gateway' subdirectory inside the 'resources' directory, and add a file named 'api_gateway_construct.py' within the 'api_gateway' directory

import aws_cdk as cdk
from aws_cdk import aws_apigatewayv2_alpha as _apigwv2
from aws_cdk import aws_apigatewayv2_authorizers_alpha as _authorizers
from aws_cdk import aws_apigatewayv2_integrations_alpha as _integration
from aws_cdk import aws_lambda as _lambda
from constructs import Construct


class APIGateway(Construct):
    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        props: dict = {},
        **kwargs
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        env = props.get('env', 'dev')
        custom_auth: _lambda.Function = props.get('custom_auth', None)
        backend_handler: _lambda.Function = props.get('backend_handler', None)

        cors_preflight = _apigwv2.CorsPreflightOptions(
            allow_origins=['*'],
            allow_methods=[_apigwv2.CorsHttpMethod.ANY],
            allow_headers=[
                'Content-Type', 'X-Amz-Date', 'X-Amz-Security-Token',
                'Authorization', 'X-Api-Key', 'X-Requested-With', 'Accept',
                'Access-Control-Allow-Methods', 'Access-Control-Allow-Origin',
                'Access-Control-Allow-Headers'
            ]
        )

        self.http_api = _apigwv2.HttpApi(
            self, f"{env}-api",
            cors_preflight=cors_preflight,
        )

        authorizer = _authorizers.HttpLambdaAuthorizer(
            'custom-authorizer-cvc', custom_auth,
            response_types=[
                _authorizers.HttpLambdaResponseType.IAM
            ],
            results_cache_ttl=cdk.Duration.hours(1)
        )

        self.integration = _integration.HttpLambdaIntegration(
            "LambdaHandler", backend_handler
        )

        self.http_api.add_routes(
            path='/admin',
            authorizer=authorizer,
            integration=self.integration,
            methods=[
                _apigwv2.HttpMethod.GET,
                _apigwv2.HttpMethod.POST,
                _apigwv2.HttpMethod.PUT,
                _apigwv2.HttpMethod.DELETE,
            ]
        )

        self.http_api.add_routes(
            path='/individual',
            authorizer=authorizer,
            integration=self.integration,
            methods=[
                _apigwv2.HttpMethod.GET,
                _apigwv2.HttpMethod.POST,
                _apigwv2.HttpMethod.PUT,
                _apigwv2.HttpMethod.DELETE,
            ]
        )

Basically, we create:

A HTTP Api
An authorizer with a lambda function that returns a policy.
Simple Lambda integration.
Routes

Build stack

Move to file role_based_access_control_stack.py, we import and create the define services

from aws_cdk import Stack
from constructs import Construct

from role_based_access_control.resources.api_gateway.api_gateway_construct import \
    APIGateway
from role_based_access_control.resources.cognito.congito_construct import \
    Cognito
from role_based_access_control.resources.dynamodb.dynamodb_construct import \
    DynamoDB
from role_based_access_control.resources.lambdas.lambda_construct import Lambda


class RoleBasedAccessControlStack(Stack):

    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        props: dict = {},
        **kwargs
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        self.props = props

        self._lambda = Lambda(
            self, 'lambda',
            props={
                **self.props
            }
        )

        self.cognito = Cognito(
            self, 'cognito',
            props={
                **self.props,
                'post_confirmation': self._lambda.post_confirmation,
                'custom_auth': self._lambda.custom_auth,
            }
        )

        self.dynamodb = DynamoDB(
            self, 'dynamodb-cvc',
            props={
                **self.props,
                'custom_auth': self._lambda.custom_auth,
            }
        )

        self.apigw = APIGateway(
            self, 'apigw',
            props={
                **self.props,
                'custom_auth': self._lambda.custom_auth,
                'backend_handler': self._lambda.backend_handler
            }
        )

and at main.py, add some props:

#!/usr/bin/env python3

import aws_cdk as cdk

from role_based_access_control.role_based_access_control_stack import \
    RoleBasedAccessControlStack

app = cdk.App()
RoleBasedAccessControlStack(
    app, "RoleBasedAccessControlStack",
    props={
        'env': 'dev',
    }
)

app.synth()

Deploy

cdk deploy

Cleanup

cdk destroy

Top comments (1)

nhipky • Sep 18 '23

It's very helpful

DEV Community