feature: add cognito identity (#126)

add cognito user pool
add testing test user
add jwt authoizers
add e2e id token generation
add security tests

---------

Co-authored-by: Ran Isenberg <ran.isenberg@ranthebuilder.cloud>

.github/workflows/pr-serverless-service.yml

@@ -108,5 +108,6 @@ jobs:
           verbose: false # optional (default = false)
       - name: Run E2E tests
         run: make e2e
-      - name: Destroy stack in non production
+      - name: Destroy stack
+        if: always()
         run: make destroy

.gitignore

@@ -253,3 +253,4 @@ lambda_requirements.txt
 
 node_modules
 .idea
+.ruff_cache

infrastructure/product/constants.py

@@ -43,3 +43,9 @@
 STREAM_TESTS_TABLE_NAME_OUTPUT = 'TestDbOutput'
 STREAM_TESTS_TABLE_NAME = 'StreamTest'
 STREAM_PROCESSOR_TEST_LAMBDA_ROLE_ARN = 'StreamTestRoleArn'
+TEST_USER_USERNAME = 'tests'
+TEST_USER_TEMP_PWD = 'DUMMYDUMaaMY1111@#22'
+TEST_USER_IDENTITY_SECRET_NAME_OUTPUT = 'TestUserSecret'
+IDENTITY_USER_POOL_ID_OUTPUT = 'UserPoolId'
+IDENTITY_USER_NAME = 'UserPool'
+IDENTITY_APP_CLIENT_ID_OUTPUT = 'AppClientId'

infrastructure/product/crud_api_construct.py

@@ -9,21 +9,24 @@
 import infrastructure.product.constants as constants
 from infrastructure.product.crud_api_db_construct import ApiDbConstruct
 from infrastructure.product.crud_monitoring import CrudMonitoring
+from infrastructure.product.identity_provider_construct import IdentityProviderConstruct
 
 
 class CrudApiConstruct(Construct):
-    def __init__(self, scope: Construct, id_: str, lambda_layer: PythonLayerVersion) -> None:
+    def __init__(self, scope: Construct, id_: str, lambda_layer: PythonLayerVersion, is_production: bool) -> None:
         super().__init__(scope, id_)
         self.api_db = ApiDbConstruct(self, f'{id_}db')
         self.common_layer = lambda_layer
+        self.idp = IdentityProviderConstruct(self, f'{id_}users', is_production)
         self.rest_api = self._build_api_gw()
         api_resource: aws_apigateway.Resource = self.rest_api.root.add_resource('api')
         product_resource = api_resource.add_resource(constants.PRODUCT_RESOURCE).add_resource('{product}')
-        self.create_prod_func = self._add_put_product_lambda_integration(product_resource, self.api_db.db, self.api_db.idempotency_db)
-        self.delete_prod_func = self._add_delete_product_lambda_integration(product_resource, self.api_db.db)
-        self.get_prod_func = self._add_get_product_lambda_integration(product_resource, self.api_db.db)
+        authorizer = aws_apigateway.CognitoUserPoolsAuthorizer(self, 'ProductsAuthorizer', cognito_user_pools=[self.idp.user_pool])
+        self.create_prod_func = self._add_put_product_lambda_integration(product_resource, self.api_db.db, self.api_db.idempotency_db, authorizer)
+        self.delete_prod_func = self._add_delete_product_lambda_integration(product_resource, self.api_db.db, authorizer)
+        self.get_prod_func = self._add_get_product_lambda_integration(product_resource, self.api_db.db, authorizer)
         products_resource: aws_apigateway.Resource = api_resource.add_resource(constants.PRODUCTS_RESOURCE)
-        self.list_prods_func = self._add_list_products_lambda_integration(products_resource, self.api_db.db)
+        self.list_prods_func = self._add_list_products_lambda_integration(products_resource, self.api_db.db, authorizer)
         # add CW dashboards
         self.dashboard = CrudMonitoring(
             self,
@@ -145,6 +148,7 @@ def _add_put_product_lambda_integration(
         put_resource: aws_apigateway.Resource,
         db: dynamodb.Table,
         idempotency_table: dynamodb.Table,
+        auth: aws_apigateway.CognitoUserPoolsAuthorizer,
     ) -> _lambda.Function:
         role = self._build_create_product_lambda_role(db, idempotency_table)
         lambda_function = _lambda.Function(
@@ -169,13 +173,19 @@ def _add_put_product_lambda_integration(
         )
 
         # PUT /api/product/{product}/
-        put_resource.add_method(http_method='PUT', integration=aws_apigateway.LambdaIntegration(handler=lambda_function))
+        put_resource.add_method(
+            http_method='PUT',
+            integration=aws_apigateway.LambdaIntegration(handler=lambda_function),
+            authorization_type=aws_apigateway.AuthorizationType.COGNITO,
+            authorizer=auth,
+        )
         return lambda_function
 
     def _add_delete_product_lambda_integration(
         self,
-        put_resource: aws_apigateway.Resource,
+        resource: aws_apigateway.Resource,
         db: dynamodb.Table,
+        auth: aws_apigateway.CognitoUserPoolsAuthorizer,
     ) -> _lambda.Function:
         role = self._build_delete_product_lambda_role(db)
         lambda_function = _lambda.Function(
@@ -199,13 +209,19 @@ def _add_delete_product_lambda_integration(
         )
 
         # DELETE /api/product/{product}/
-        put_resource.add_method(http_method='DELETE', integration=aws_apigateway.LambdaIntegration(handler=lambda_function))
+        resource.add_method(
+            http_method='DELETE',
+            integration=aws_apigateway.LambdaIntegration(handler=lambda_function),
+            authorization_type=aws_apigateway.AuthorizationType.COGNITO,
+            authorizer=auth,
+        )
         return lambda_function
 
     def _add_get_product_lambda_integration(
         self,
-        put_resource: aws_apigateway.Resource,
+        resource: aws_apigateway.Resource,
         db: dynamodb.Table,
+        auth: aws_apigateway.CognitoUserPoolsAuthorizer,
     ) -> _lambda.Function:
         role = self._build_get_product_lambda_role(db)
         lambda_function = _lambda.Function(
@@ -229,13 +245,19 @@ def _add_get_product_lambda_integration(
         )
 
         # GET /api/product/{product}/
-        put_resource.add_method(http_method='GET', integration=aws_apigateway.LambdaIntegration(handler=lambda_function))
+        resource.add_method(
+            http_method='GET',
+            integration=aws_apigateway.LambdaIntegration(handler=lambda_function),
+            authorization_type=aws_apigateway.AuthorizationType.COGNITO,
+            authorizer=auth,
+        )
         return lambda_function
 
     def _add_list_products_lambda_integration(
         self,
         api_resource: aws_apigateway.Resource,
         db: dynamodb.Table,
+        auth: aws_apigateway.CognitoUserPoolsAuthorizer,
     ) -> _lambda.Function:
         role = self._build_list_products_lambda_role(db)
         lambda_function = _lambda.Function(
@@ -259,6 +281,11 @@ def _add_list_products_lambda_integration(
         )
 
         # GET /api/products/
-        api_resource.add_method(http_method='GET', integration=aws_apigateway.LambdaIntegration(handler=lambda_function))
+        api_resource.add_method(
+            http_method='GET',
+            integration=aws_apigateway.LambdaIntegration(handler=lambda_function),
+            authorization_type=aws_apigateway.AuthorizationType.COGNITO,
+            authorizer=auth,
+        )
 
         return lambda_function

infrastructure/product/custom_resource_handler.py

@@ -0,0 +1,24 @@
+import json
+
+import boto3
+
+
+def handler(event, context):
+    if event['RequestType'] == 'Create' or event['RequestType'] == 'Update':
+        secret_name = event['ResourceProperties']['SecretName']
+        user_pool_id = event['ResourceProperties']['UserPoolId']
+
+        # Fetch the secret value
+        secrets_client = boto3.client('secretsmanager')
+        response = secrets_client.get_secret_value(SecretId=secret_name)
+        secret_value = json.loads(response['SecretString'])
+
+        # Change Cognito user password
+        cognito_client = boto3.client('cognito-idp')
+        cognito_client.admin_set_user_password(
+            UserPoolId=user_pool_id,
+            Username=secret_value['username'],
+            Password=secret_value['password'],
+            Permanent=True,
+        )
+        return {'PhysicalResourceId': 'AWS:CustomSetPass'}

infrastructure/product/identity_provider_construct.py

@@ -0,0 +1,134 @@
+from os import path
+
+from aws_cdk import CfnOutput, CustomResource, Duration, Fn, RemovalPolicy
+from aws_cdk import aws_cognito as cognito
+from aws_cdk import aws_iam as iam
+from aws_cdk import aws_lambda as _lambda
+from aws_cdk import aws_secretsmanager as secrets
+from aws_cdk.custom_resources import AwsCustomResource, AwsCustomResourcePolicy, PhysicalResourceId, Provider
+from constructs import Construct
+
+import infrastructure.product.constants as constants
+
+
+class IdentityProviderConstruct(Construct):
+    def __init__(self, scope: Construct, id_: str, is_production: bool) -> None:
+        super().__init__(scope, id_)
+        self.id_ = id_
+        self.user_pool = self._create_user_pool(is_production)
+        self.app_client = self._create_app_client(self.user_pool)
+        if not is_production:
+            self._create_test_user(self.user_pool)
+
+    def _create_app_client(self, user_pool: cognito.UserPool) -> cognito.UserPoolClient:
+        app_client = cognito.UserPoolClient(
+            self,
+            'UserPoolClient',
+            user_pool=user_pool,
+            auth_flows=cognito.AuthFlow(user_password=True),  # Allow username/password authentication
+        )
+
+        CfnOutput(self, id=constants.IDENTITY_APP_CLIENT_ID_OUTPUT, value=app_client.user_pool_client_id).override_logical_id(
+            constants.IDENTITY_APP_CLIENT_ID_OUTPUT
+        )
+        return app_client
+
+    def _create_user_pool(self, is_production: bool) -> cognito.UserPool:
+        user_pool = cognito.UserPool(
+            self,
+            'UserPool',
+            user_pool_name=constants.IDENTITY_USER_NAME,
+            advanced_security_mode=cognito.AdvancedSecurityMode.ENFORCED,
+            password_policy=cognito.PasswordPolicy(
+                min_length=12,
+                require_lowercase=True,
+                require_uppercase=True,
+                require_digits=True,
+                require_symbols=True,
+                temp_password_validity=Duration.days(7),
+            ),
+            sign_in_aliases=cognito.SignInAliases(username=True, email=True),
+            removal_policy=RemovalPolicy.DESTROY,
+            mfa=cognito.Mfa.OPTIONAL if not is_production else cognito.Mfa.REQUIRED,
+        )
+
+        CfnOutput(self, id=constants.IDENTITY_USER_POOL_ID_OUTPUT, value=user_pool.user_pool_id).override_logical_id(
+            constants.IDENTITY_USER_POOL_ID_OUTPUT
+        )
+        return user_pool
+
+    def _create_test_user(self, user_pool: cognito.UserPool):
+        create_test_user = AwsCustomResource(
+            self,
+            'AwsCustom-CreateUser',
+            on_create={
+                'service': 'CognitoIdentityServiceProvider',
+                'action': 'adminCreateUser',
+                'parameters': {
+                    'UserPoolId': user_pool.user_pool_id,
+                    'Username': constants.TEST_USER_USERNAME,
+                    'MessageAction': 'SUPPRESS',
+                    'TemporaryPassword': constants.TEST_USER_TEMP_PWD,
+                },
+                'physical_resource_id': PhysicalResourceId.of(f'{self.id_}CreateTestUserResource'),
+            },
+            policy=AwsCustomResourcePolicy.from_sdk_calls(resources=[user_pool.user_pool_arn]),
+            install_latest_aws_sdk=True,
+            removal_policy=RemovalPolicy.DESTROY,
+            timeout=Duration.minutes(1),
+        )
+        create_test_user.node.add_dependency(user_pool.node.default_child)
+
+        # Save the user's credentials to AWS Secrets Manager
+        user_credentials_secret = secrets.Secret(
+            self,
+            f'{self.id_}TestUserSecret',
+            description='Credentials for the test user in the Cognito User Pool',
+            generate_secret_string=secrets.SecretStringGenerator(
+                secret_string_template=Fn.sub('{"username": "${user}"}', {'user': constants.TEST_USER_USERNAME}),
+                generate_string_key='password',
+                exclude_punctuation=False,
+                exclude_lowercase=False,
+                exclude_numbers=False,
+                exclude_uppercase=False,
+                include_space=False,
+                password_length=16,
+                require_each_included_type=True,
+            ),
+        )
+
+        user_credentials_secret.node.add_dependency(create_test_user)
+
+        CfnOutput(self, id=constants.TEST_USER_IDENTITY_SECRET_NAME_OUTPUT, value=user_credentials_secret.secret_name).override_logical_id(
+            constants.TEST_USER_IDENTITY_SECRET_NAME_OUTPUT
+        )
+
+        with open(path.join(path.dirname(__file__), 'custom_resource_handler.py'), encoding='utf-8') as file:
+            handler_code = file.read()
+
+        lambda_function = _lambda.Function(
+            self,
+            'SetPassHandler',
+            runtime=_lambda.Runtime.PYTHON_3_11,
+            handler='index.handler',
+            code=_lambda.InlineCode(handler_code),
+        )
+
+        # Grant permission to Lambda to get secrets and interact with Cognito
+        user_credentials_secret.grant_read(lambda_function)
+        lambda_function.add_to_role_policy(iam.PolicyStatement(actions=['cognito-idp:AdminSetUserPassword'], resources=[user_pool.user_pool_arn]))
+
+        provider = Provider(scope=self, id=f'{self.id_}Provider', on_event_handler=lambda_function)
+        lambda_resource = CustomResource(
+            scope=self,
+            id=f'{self.id_}SetAdminFunc',
+            service_token=provider.service_token,
+            removal_policy=RemovalPolicy.DESTROY,
+            properties={
+                'SecretName': user_credentials_secret.secret_name,
+                'UserPoolId': user_pool.user_pool_id,
+            },
+            resource_type='Custom::IdentitySetTestPassword',
+        )
+
+        lambda_resource.node.add_dependency(user_credentials_secret)

infrastructure/product/service_stack.py

@@ -21,6 +21,7 @@ def __init__(self, scope: Construct, id: str, is_production: bool, **kwargs) ->
             self,
             id_=get_construct_name(id, constants.CRUD_CONSTRUCT_NAME),
             lambda_layer=self.shared_layer,
+            is_production=is_production,
         )
 
         self.stream_processor = StreamProcessorConstruct(
@@ -70,5 +71,6 @@ def _add_security_tests(self) -> None:
                 {'id': 'AwsSolutions-APIG6', 'reason': 'not mandatory in a sample template'},
                 {'id': 'AwsSolutions-APIG4', 'reason': 'authorization not mandatory in a sample template'},
                 {'id': 'AwsSolutions-COG4', 'reason': 'not using cognito'},
+                {'id': 'AwsSolutions-SMG4', 'reason': 'secret for cognito does not support auto rotate'},
             ],
         )

infrastructure/product/stream_processor_testing_construct.py

@@ -77,6 +77,7 @@ def _build_state_machine(self, lambda_func: _lambda.Function, event_bus: events.
                     self,
                     f'{self.id_}TestStateMachineLogGroup',
                     removal_policy=RemovalPolicy.DESTROY,
+                    log_group_name=f'/aws/vendedlogs/states/{self.id_}',
                 ),
                 level=sfn.LogLevel.ALL,
                 include_execution_data=True,

tests/e2e/crud/conftest.py

@@ -1,8 +1,17 @@
+import json
 from typing import Generator
 
+import boto3
 import pytest
 
-from infrastructure.product.constants import APIGATEWAY, PRODUCT_RESOURCE, PRODUCTS_RESOURCE, TABLE_NAME_OUTPUT
+from infrastructure.product.constants import (
+    APIGATEWAY,
+    IDENTITY_APP_CLIENT_ID_OUTPUT,
+    PRODUCT_RESOURCE,
+    PRODUCTS_RESOURCE,
+    TABLE_NAME_OUTPUT,
+    TEST_USER_IDENTITY_SECRET_NAME_OUTPUT,
+)
 from product.crud.models.product import Product
 from tests.crud_utils import clear_table, generate_product_id
 from tests.e2e.crud.utils import create_product, delete_product
@@ -35,10 +44,35 @@ def table_name():
 
 
 @pytest.fixture()
-def add_product_entry_to_db(api_gw_url_slash_product: str, table_name: str) -> Generator[Product, None, None]:
+def add_product_entry_to_db(api_gw_url_slash_product: str, table_name: str, id_token: str) -> Generator[Product, None, None]:
     clear_table(table_name)
     product_id = generate_product_id()
     product = Product(id=product_id, price=1, name='test')
-    create_product(api_gw_url_slash_product=api_gw_url_slash_product, product_id=product_id, price=product.price, name=product.name)
+    create_product(
+        api_gw_url_slash_product=api_gw_url_slash_product, product_id=product_id, price=product.price, name=product.name, id_token=id_token
+    )
     yield product
-    delete_product(api_gw_url_slash_product, product_id)
+    delete_product(api_gw_url_slash_product, product_id, id_token)
+
+
+@pytest.fixture(scope='session', autouse=True)
+def id_token() -> str:
+    # Initialize boto3 client for Secrets Manager and Cognito Identity
+    secrets_manager_client = boto3.client('secretsmanager')
+    cognito_client = boto3.client('cognito-idp')
+
+    # Name of the secret created in the previous CDK stack
+    SECRET_NAME = get_stack_output(TEST_USER_IDENTITY_SECRET_NAME_OUTPUT)
+    CLIENT_ID = get_stack_output(IDENTITY_APP_CLIENT_ID_OUTPUT)
+
+    # Retrieve secret from Secrets Manager
+    response = secrets_manager_client.get_secret_value(SecretId=SECRET_NAME)
+    secret = json.loads(response['SecretString'])
+
+    # Use the credentials from the secret to login to the Cognito User Pool
+    auth_response = cognito_client.initiate_auth(
+        AuthFlow='USER_PASSWORD_AUTH',
+        AuthParameters={'USERNAME': secret['username'], 'PASSWORD': secret['password']},
+        ClientId=CLIENT_ID,
+    )
+    return auth_response['AuthenticationResult']['IdToken']