Role-based access control with AWS CDK
Tags: rbac, cdk, aws, python
Published: 2023-09-17 08:10:50

## 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.


![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/xujdyty90ey7jm3m0ho8.png)

The GitHub repo is [here](https://github.com/toantranct94/role-based-access-control)

## 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
```