diff --git a/apigw-bedrock-cognito-cdk/.gitignore b/apigw-bedrock-cognito-cdk/.gitignore
new file mode 100644
index 000000000..37833f8be
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/.gitignore
@@ -0,0 +1,10 @@
+*.swp
+package-lock.json
+__pycache__
+.pytest_cache
+.venv
+*.egg-info
+
+# CDK asset staging directory
+.cdk.staging
+cdk.out
diff --git a/apigw-bedrock-cognito-cdk/README.md b/apigw-bedrock-cognito-cdk/README.md
new file mode 100644
index 000000000..678358e10
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/README.md
@@ -0,0 +1,224 @@
+# Access Bedrock via API Gateway with Cognito User Management
+
+This pattern creates an AWS CDK Python application to access Bedrock via API Gateway with Cognito user management, domain restriction, API request throttling, and quota limits.
+
+Learn more about this pattern at Serverless Land Patterns: << Add the live URL here >>
+
+Important: this application uses various AWS services and there are costs associated with these services after the Free Tier usage - please see the [AWS Pricing page](https://aws.amazon.com/pricing/) for details. You are responsible for any AWS costs incurred. No warranty is implied in this example.
+
+## Requirements
+
+* [Create an AWS account](https://portal.aws.amazon.com/gp/aws/developer/registration/index.html) if you do not already have one and log in. The IAM user that you use must have sufficient permissions to make necessary AWS service calls and manage AWS resources.
+* [Install and Configure AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/install-cliv2.html)
+* [Install Git](https://git-scm.com/book/en/v2/Getting-Started-Installing-Git)
+* [Install Node and NPM](https://nodejs.org/en/download/)
+* [Install AWS Cloud Development Kit (AWS CDK)](https://docs.aws.amazon.com/cdk/latest/guide/cli.html)
+* [Install Python 3](https://www.python.org/downloads/)
+* [Install Docker](https://www.docker.com/products/docker-desktop/)
+* [Grant Bedrock Model Access](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html)
+
+## Deployment Instructions
+
+1. Create a new directory, navigate to that directory in a terminal and clone the GitHub repository:
+ ```
+ git clone https://github.com/aws-samples/serverless-patterns
+ ```
+2. Change directory to the pattern directory:
+ ```
+ cd apigw-bedrock-cognito-cdk
+ ```
+3. Create a virtual environment for Python:
+ ```
+ python3 -m venv .venv
+ ```
+4. Activate the virtual environment:
+ ```
+ source .venv/bin/activate
+ ```
+ For a Windows platform, activate the virtualenv like this:
+ ```
+ .venv\Scripts\activate.bat
+ ```
+5. Install the required Python dependencies:
+ ```
+ pip install -r requirements.txt
+ ```
+6. Bootstrap the AWS environment, if you haven't already done so:
+ ```
+ cdk bootstrap
+ ```
+7. Review the CloudFormation template AWS CDK generates for the stack:
+ ```
+ cdk synth
+ ```
+8. Deploy the AWS resources:
+ ```
+ cdk deploy
+ ```
+ Optionally you can add the optional variables using the `--context`, see [API Gateway Configuration](#API-Gateway-Configuration) and [Cognito Integration and Configuration](#Cognito-Integration-and-Configuration).
+ ```
+ cdk deploy --context API_THROTTLE_RATE_LIMIT=1 --context API_THROTTLE_BURST_LIMIT=2 --context API_QUOTA_LIMIT=25 --context API_QUOTA_PERIOD=DAY --context ORGANIZATION_DOMAIN=@example.com
+ ```
+9. Note the outputs from the CDK deployment process. These contain the resource names and/or ARNs which are used for testing.
+
+## How it works
+
+### Overview
+
+This pattern deploys an Amazon API Gateway REST API with the following routes: `POST /register`, `POST /login`, `GET` and `POST /bedrock`. It includes a Amazon Cognito User Pool and Lambdas to handle requests from the API Gateway. The API Gateway allows CORS for all origins and methods, incorporates an Usage Plan, and has throttle and quota limits for the `/bedrock` endpoint. The `/bedrock` endpoint allows access to Amazon Bedrock Foundation models.
+
+### Components and Configuration
+
+#### API Gateway Routes
+- `/register` endpoint: Accepts `email`, `password`, and `fullname` in the body, interacts with a proxy lambda integration to register users to the Cognito User Pool. Returns an API Key which will be associated with the API Gateway's usage plan and stored within the Cognito user's custom `api_key` field . If an organization domain is specified in the `ORGANIZATION_DOMAIN` context variable, a pre signup lambda is provisioned to reject users not belonging to the specified domain.
+
+- `/login` endpoint: Accepts `email` and `password` in the body, interacts with a proxy lambda integration to authenticate the user. Returns an bearer token, containing `IdToken`, `AccessToken`, `RefreshToken` and other metadata. If the user loses their API key, they can decrypt the `IdToken` using [jwt.io](https://jwt.io/) or other libraries to retrieve the API key from the `custom:api_key` field.
+
+- `/bedrock` endpoint: Protected with a Cognito authorizer to ensure only requests with valid `Authorization` and `x-api-key` tokens in headers can access the endpoint, interacts with a proxy lambda integration. A `GET` request lists all foundation models, and a `POST` request takes `modelId` and `inferenceParameters` in the body, to invoke and return response from the foundation model.
+
+#### API Gateway Configuration
+- CORS: Enabled for all origins and methods.
+- Usage Plan: Configured to manage API access.
+- Throttle Limits: Rate limit of 1 request per second with a burst limit of 2 requests.
+- Quota Limits: Set to 25 requests per day for the `/bedrock` endpoint.
+- These limits can be modified during deployment using context variables (`API_THROTTLE_RATE_LIMIT`, `API_THROTTLE_BURST_LIMIT`, `API_QUOTA_LIMIT`, `API_QUOTA_PERIOD`)
+- Logging: Enabled for all Error and Info request logs.
+
+#### Cognito Integration and Configuration
+- User Pool: Manages user registration and login.
+- Organization Domain Restriction: The organization domain restriction can be adjusted during deployment using the context variable `ORGANIZATION_DOMAIN`. A Pre SignUp Lambda trigger will be added to enforce specific domain restrictions.
+
+#### Lambda Integration
+- `bedrock.py`:
+ - Uses the `boto3` library to make API calls to Amazon Bedrock APIs.
+ - Utilizes the `jsonpath_ng` library to dynamically map and retrieve responses from foundation models provided by Anthropic, AI21 Labs, Amazon, and Meta.
+- `auth.py`:
+ - Uses the `boto3` library to make API calls to Amazon Cognito and Amazon API Gateway.
+ - Manages user creation, login, and the creation and association of API keys.
+- `pre_signup.py` (if valid):
+ - Validates user email domain during registration.
+- All Lambda Configuration:
+ - Timeout: Set to 29 second due to maximum integration timeout limit - [Amazon API Gateway Limits](https://docs.aws.amazon.com/apigateway/latest/developerguide/limits.html).
+ - Logging: All events are logged to Amazon CloudWatch, with a custom redaction function to remove passwords from the `auth.py` Lambda prior to logging.
+
+### Request and Response Examples
+
+- **Register User**
+ - **Request Body:**
+ ```json
+ {
+ "email": "user@example.com",
+ "password": "securePassword123",
+ "fullname": "John Doe"
+ }
+ ```
+ - **Response Body:**
+ ```json
+ {
+ "status": 200,
+ "message": "User user@example.com created successfully.",
+ "data": {"API Key": "generatedApiKey"},
+ "success": true,
+ }
+ ```
+
+- **Login User**
+ - **Request Body:**
+ ```json
+ {
+ "email": "user@example.com",
+ "password": "securePassword123"
+ }
+ ```
+ - **Response Body:**
+ ```json
+ {
+ "status": 200,
+ "IdToken": "generatedIdToken",
+ "AccessToken": "generatedAccessToken",
+ "RefreshToken": "generatedRefreshToken",
+ "ExpiresIn": 3600,
+ "TokenType": "Bearer",
+ "sucess": true
+ }
+ ```
+
+- **Make Bedrock Request**
+ - **Request Headers:**
+ - Authorization: Bearer [IdToken]
+ - x-api-key: [APIKey]
+ - **GET /bedrock Response Body:**
+ ```json
+ {
+ "status": 200,
+ "foundationModels": [...],
+ "message": "Successfully retrieved foundation models list."
+ }
+ ```
+ - **POST /bedrock Request Body:**
+ ```json
+ {
+ "modelId": "exampleModelId",
+ "inferenceParameters": {...}
+ }
+ ```
+ - **POST /bedrock Response Body:**
+ ```json
+ {
+ "status": 200,
+ "message": "Successfully retrieved response from foundation model.",
+ "data": "..."
+ }
+ ```
+
+## Testing
+
+Prior to running the tests, ensure that your account has the necessary access to the following Bedrock models: `ai21.j2-mid-v1`, `anthropic.claude-instant-v1`, `amazon.titan-text-express-v1`, and `meta.llama2-13b-chat-v1`. Follow the guide provided in the [Requirements](#Requirements) - [Bedrock model access granted](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html) to grant access to these models. Without access, the tests will fail due to insufficient permissions to interact with the foundation models.
+
+1. Activate the virtual environment, if you haven't already done so from the deployment instructions:
+ ```
+ source .venv/bin/activate
+ ```
+ For a Windows platform, activate the virtualenv like this:
+ ```
+ .venv\Scripts\activate.bat
+ ```
+2. Install the Python required dependencies:
+ ```
+ pip install -r requirements-dev.txt
+ ```
+3. From the CDK output during the deployment process, obtain the values for Cognito User Pool ID and Rest API Endpoint:
+ ```
+ ApigwBedrockCognitoCdkStack.CognitoUserPoolID = us-east-1_XXXXXXXXX
+ ApigwBedrockCognitoCdkStack.RestAPIEndpoint = https://XXXXXXXXX.execute-api.us-east-1.amazonaws.com/prod/
+ ```
+4. Update the `API_ENDPOINT` and `USER_POOL_ID` variables in `tests/e2e/test_apigw_bedrock_cognito_cdk_stack.py` with the obtained values:
+ ```
+ API_ENDPOINT = "https://XXXXXXXXX.execute-api.us-east-1.amazonaws.com/prod/"
+ USER_POOL_ID = "us-east-1_XXXXXXXXX"
+ ```
+ If organization domain restriction is enabled, modify the `TEST_EMAIL` to be an email with your domain, and `NON_ORG_TEST_EMAIL` to an email without your domain:
+ ```
+ TEST_EMAIL = "johndoe@org.com"
+ NON_ORG_TEST_EMAIL = "johndoe@example.com"
+ ```
+5. Execute Pytest:
+ ```
+ pytest -v
+ ```
+
+## Cleanup
+
+1. Delete the stack:
+ ```
+ cdk destroy
+ ```
+2. Delete all API Keys:
+ Before executing the following command, be aware that it will delete all API keys in the account. Ensure you have the necessary backups or are certain of the consequences.
+ ```
+ sh utils/delete_all_api_keys.sh
+ ```
+----
+Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+SPDX-License-Identifier: MIT-0
\ No newline at end of file
diff --git a/apigw-bedrock-cognito-cdk/apigw-bedrock-cognito-cdk.json b/apigw-bedrock-cognito-cdk/apigw-bedrock-cognito-cdk.json
new file mode 100644
index 000000000..ce0619e98
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/apigw-bedrock-cognito-cdk.json
@@ -0,0 +1,109 @@
+{
+ "title": "API Gateway with Cognito to Lambda to Amazon Bedrock",
+ "description": "Create an AWS CDK Python application for Bedrock API with Cognito user management, domain restriction, API request throttling, and quota limits.",
+ "language": "Python",
+ "level": "200",
+ "framework": "CDK",
+ "introBox": {
+ "headline": "How it works",
+ "text": [
+ "This sample project demonstrates how to create an API Gateway with AWS CDK, backed by a Cognito user pool for user management.",
+ "Users can register and log in via an endpoint. Registration can be restricted to a specific domain, such as '@example.com'. The application supports request throttling and quota limits to the API.",
+ "Configuration for domain restriction, request throttling, and quota management can be done using CDK context variables."
+ ]
+ },
+ "gitHub": {
+ "template": {
+ "repoURL": "https://github.com/aws-samples/serverless-patterns/tree/main/apigw-bedrock-cognito-cdk",
+ "templateURL": "serverless-patterns/apigw-bedrock-cognito-cdk",
+ "projectFolder": "apigw-bedrock-cognito-cdk",
+ "templateFile": "app.py"
+ }
+ },
+ "resources": {
+ "bullets": [
+ {
+ "text": "Amazon API Gateway - REST API",
+ "link": "https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-rest-api.html"
+ },
+ {
+ "text": "Amazon Bedrock",
+ "link": "https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html"
+ },
+ {
+ "text": "Amazon Bedrock - Inference parameters for foundation models",
+ "link": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html"
+ },
+ {
+ "text": "AWS CDK v2 - Python",
+ "link": "https://docs.aws.amazon.com/cdk/v2/guide/work-with-cdk-python.html"
+ },
+ {
+ "text": "Amazon Cognito - User Pools",
+ "link": "https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-identity-pools.html"
+ }
+ ]
+ },
+ "deploy": {
+ "text": [
+ "Deploy the stack: cdk deploy
",
+ "Deploy the stack with context variables: cdk deploy --context ORGANIZATION_DOMAIN=@example.com --context API_THROTTLE_RATE_LIMIT=1 --context API_THROTTLE_BURST_LIMIT=2 --context API_QUOTA_LIMIT=25 --contect API_QUOTA_LIMIT=DAY
"
+ ]
+ },
+ "testing": {
+ "text": [
+ "Run end-to-end tests: pytest
"
+ ]
+ },
+ "cleanup": {
+ "text": [
+ "Delete the stack: cdk destroy
."
+ ]
+ },
+ "authors": [
+ {
+ "name": "Inthuson Anandakumar",
+ "image": "https://avatars.githubusercontent.com/u/8601106?v=4",
+ "bio": "Inthuson Anandakumar Technical Account Manager at AWS based out of London. Inthu currently works with strategic financial services customers for AWS, helping them architect, build, and optimize cloud applications. With a background in software development and solution architecture, he thrives on crafting innovative solutions and enjoys sharing his discoveries with customers across a wide array of sectors.",
+ "linkedin": "inthuson"
+ }
+ ],
+ "patternArch": {
+ "icon1": {
+ "x": 10,
+ "y": 50,
+ "service": "cognito",
+ "label": "Amazon Cognito"
+ },
+ "icon2": {
+ "x": 38,
+ "y": 50,
+ "service": "apigw",
+ "label": "Amazon API Gateway"
+ },
+ "icon3": {
+ "x": 66,
+ "y": 50,
+ "service": "lambda",
+ "label": "AWS Lambda"
+ },
+ "icon4": {
+ "x": 90,
+ "y": 50,
+ "service": "bedrock",
+ "label": "Amazon Bedrock"
+ },
+ "line1": {
+ "from": "icon1",
+ "to": "icon2"
+ },
+ "line2": {
+ "from": "icon2",
+ "to": "icon3"
+ },
+ "line3": {
+ "from": "icon3",
+ "to": "icon4"
+ }
+ }
+}
diff --git a/apigw-bedrock-cognito-cdk/apigw_bedrock_cognito_cdk/__init__.py b/apigw-bedrock-cognito-cdk/apigw_bedrock_cognito_cdk/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/apigw-bedrock-cognito-cdk/apigw_bedrock_cognito_cdk/apigw_bedrock_cognito_cdk_stack.py b/apigw-bedrock-cognito-cdk/apigw_bedrock_cognito_cdk/apigw_bedrock_cognito_cdk_stack.py
new file mode 100644
index 000000000..995fbbac5
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/apigw_bedrock_cognito_cdk/apigw_bedrock_cognito_cdk_stack.py
@@ -0,0 +1,250 @@
+from aws_cdk import (
+ Stack,
+ aws_cognito as cognito,
+ RemovalPolicy,
+ aws_lambda as _lambda,
+ aws_apigateway as apigateway,
+ aws_iam as iam,
+ aws_ssm as ssm,
+ CfnOutput,
+ Duration,
+ BundlingOptions,
+ Aws as aws,
+)
+from constructs import Construct
+import os
+
+
+class ApigwBedrockCognitoCdkStack(Stack):
+ def __init__(
+ self,
+ scope: Construct,
+ construct_id: str,
+ organization_domain: str,
+ api_throttle_settings: dict,
+ api_quota_settings: dict,
+ stage_name: str = "prod",
+ **kwargs,
+ ) -> None:
+ super().__init__(scope, construct_id, **kwargs)
+
+ # Create and configure Cognito User Pool
+ self.user_pool = self.create_user_pool()
+ self.ssm_user_pool_id = self.store_in_parameter_store(
+ f"/{construct_id}/user_pool/id", self.user_pool.user_pool_id, "UserPoolID"
+ )
+
+ # Output Cognito User Pool ID
+ CfnOutput(self, "CognitoUserPoolID", value=self.user_pool.user_pool_id)
+
+ # Create and configure Cognito User Pool Client
+ self.user_pool_client = self.create_user_pool_client(self.user_pool)
+ self.ssm_user_pool_client_id = self.store_in_parameter_store(
+ f"/{construct_id}/user_pool/client_id",
+ self.user_pool_client.user_pool_client_id,
+ "UserPoolClientID",
+ )
+
+ # Create and configure REST API and Usage Plan
+ self.rest_api = self.create_rest_api(stage_name)
+ api_throttle_settings = self.create_throttle_constructor(api_throttle_settings)
+ api_quota_settings = self.create_quota_constructor(api_quota_settings)
+ self.api_usage_plan = self.create_usage_plan(
+ api_throttle_settings, api_quota_settings
+ )
+ self.ssm_api_usage_plan_id = self.store_in_parameter_store(
+ f"/{construct_id}/api/usage_plan/id",
+ self.api_usage_plan.usage_plan_id,
+ "APIUsagePlanID",
+ )
+
+ # Configure pre-sign-up Lambda function
+ if organization_domain:
+ self.create_pre_signup_lambda(organization_domain)
+
+ # Create and configure Lambda function for authentication
+ self.auth_lambda = self.create_lambda_function(
+ "AuthFunc", "auth", environment={"CONSTRUCT_ID": construct_id}
+ )
+ self.add_permissions_to_lambda(
+ self.auth_lambda,
+ [
+ "apigateway:POST",
+ "ssm:GetParameter",
+ "cognito-idp:AdminUpdateUserAttributes",
+ ],
+ [
+ "*",
+ f"arn:aws:ssm:{aws.REGION}:{aws.ACCOUNT_ID}:parameter/{construct_id}/*",
+ f"arn:aws:cognito-idp:{aws.REGION}:{aws.ACCOUNT_ID}:userpool/{self.user_pool.user_pool_id}",
+ ],
+ )
+
+ # Create and configure Lambda function for Bedrock interaction
+ self.bedrock_lambda = self.create_lambda_function(
+ "BedrockFunc",
+ "bedrock",
+ None,
+ BundlingOptions(
+ image=_lambda.Runtime.PYTHON_3_12.bundling_image,
+ command=[
+ "bash",
+ "-c",
+ "pip install --no-cache -r requirements.txt -t /asset-output && cp -au . /asset-output",
+ ],
+ ),
+ )
+ self.add_permissions_to_lambda(
+ self.bedrock_lambda,
+ ["bedrock:ListFoundationModels", "bedrock:InvokeModel"],
+ ["*"],
+ )
+
+ # Create API resources
+ self.create_api_resources()
+
+ # Method to create Cognito User Pool
+ def create_user_pool(self):
+ return cognito.UserPool(
+ self,
+ "UserPool",
+ user_pool_name="UserPool",
+ self_sign_up_enabled=True,
+ sign_in_case_sensitive=False,
+ mfa=cognito.Mfa.OFF,
+ standard_attributes=cognito.StandardAttributes(
+ fullname=cognito.StandardAttribute(required=True, mutable=True),
+ email=cognito.StandardAttribute(required=True, mutable=True),
+ ),
+ custom_attributes={"api_key": cognito.StringAttribute(mutable=True)},
+ removal_policy=RemovalPolicy.DESTROY,
+ )
+
+ # Method to create Cognito User Pool Client
+ def create_user_pool_client(self, user_pool: cognito.UserPool):
+ return cognito.UserPoolClient(
+ self,
+ "UserPoolClient",
+ user_pool=user_pool,
+ auth_flows=cognito.AuthFlow(
+ user_password=True,
+ ),
+ refresh_token_validity=Duration.hours(1),
+ )
+
+ # Method to create pre-sign-up Lambda trigger
+ def create_pre_signup_lambda(self, organization_domain):
+ user_pool = self.user_pool
+ user_pool.add_trigger(
+ cognito.UserPoolOperation.PRE_SIGN_UP,
+ self.create_lambda_function(
+ "PreSignUpFunc",
+ "pre_signup",
+ {"ORGANIZATION_DOMAIN": organization_domain},
+ ),
+ )
+
+ # Method to create Lambda function
+ def create_lambda_function(
+ self, func_id: str, handler: str, environment=None, bundle_options=None
+ ):
+ return _lambda.Function(
+ self,
+ func_id,
+ runtime=_lambda.Runtime.PYTHON_3_12,
+ code=_lambda.Code.from_asset(
+ os.path.join("src", handler), bundling=bundle_options
+ ),
+ handler=f"{handler}.handler",
+ environment=environment or {},
+ timeout=Duration.seconds(29),
+ )
+
+ # Method to add permissions to Lambda function
+ def add_permissions_to_lambda(self, lambda_function, actions, resources):
+ lambda_function.add_to_role_policy(
+ iam.PolicyStatement(
+ effect=iam.Effect.ALLOW,
+ actions=actions,
+ resources=resources,
+ ),
+ )
+
+ # Method to create REST API
+ def create_rest_api(self, stage_name):
+ return apigateway.RestApi(
+ self,
+ "RestAPI",
+ default_cors_preflight_options={
+ "allow_origins": apigateway.Cors.ALL_ORIGINS,
+ "allow_methods": apigateway.Cors.ALL_METHODS,
+ },
+ cloud_watch_role=True,
+ cloud_watch_role_removal_policy=RemovalPolicy.DESTROY,
+ deploy_options=apigateway.StageOptions(
+ logging_level=apigateway.MethodLoggingLevel.INFO,
+ stage_name=stage_name,
+ ),
+ )
+
+ # Method to create API resources
+ def create_api_resources(self):
+ authorizer = apigateway.CognitoUserPoolsAuthorizer(
+ self, "CognitoAuthorizer", cognito_user_pools=[self.user_pool]
+ )
+
+ register = self.rest_api.root.add_resource("register")
+ register.add_method("POST", apigateway.LambdaIntegration(self.auth_lambda))
+
+ login = self.rest_api.root.add_resource("login")
+ login.add_method("POST", apigateway.LambdaIntegration(self.auth_lambda))
+
+ bedrock = self.rest_api.root.add_resource("bedrock")
+ bedrock.add_method(
+ "ANY",
+ apigateway.LambdaIntegration(self.bedrock_lambda),
+ authorization_type=apigateway.AuthorizationType.COGNITO,
+ authorizer=authorizer,
+ api_key_required=True,
+ )
+
+ # Method to create API throttle settings
+ def create_throttle_constructor(self, config: dict()):
+ return apigateway.ThrottleSettings(
+ rate_limit=config.get("rate_limit"),
+ burst_limit=config.get("burst_limit"),
+ )
+
+ # Method to create API quota settings
+ def create_quota_constructor(self, config: dict()):
+ period = config.get("period").upper()
+ if period == "DAY":
+ period = apigateway.Period.DAY
+ elif period == "WEEK":
+ period = apigateway.Period.WEEK
+ else:
+ period = apigateway.Period.MONTH
+
+ return apigateway.QuotaSettings(
+ limit=config.get("limit"),
+ period=period,
+ )
+
+ # Method to create API usage plan
+ def create_usage_plan(
+ self,
+ throttle_settings: apigateway.ThrottleSettings = None,
+ quota_settings: apigateway.QuotaSettings = None,
+ ):
+ return self.rest_api.add_usage_plan(
+ "APIUsagePlan",
+ api_stages=[
+ apigateway.UsagePlanPerApiStage(stage=self.rest_api.deployment_stage)
+ ],
+ throttle=throttle_settings,
+ quota=quota_settings,
+ )
+
+ # Method to store values in AWS Systems Manager Parameter Store
+ def store_in_parameter_store(self, name, value, id):
+ return ssm.StringParameter(self, id, parameter_name=name, string_value=value)
diff --git a/apigw-bedrock-cognito-cdk/app.py b/apigw-bedrock-cognito-cdk/app.py
new file mode 100644
index 000000000..749d97f8c
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/app.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+import os
+
+import aws_cdk as cdk
+
+from apigw_bedrock_cognito_cdk.apigw_bedrock_cognito_cdk_stack import (
+ ApigwBedrockCognitoCdkStack,
+)
+
+
+app = cdk.App()
+
+organization_domain = app.node.try_get_context("ORGANIZATION_DOMAIN")
+api_throttle_rate_limit = app.node.try_get_context("API_THROTTLE_RATE_LIMIT") or 1
+api_throttle_burst_limit = app.node.try_get_context("API_THROTTLE_BURST_LIMIT") or 2
+api_throttle_settings = {
+ "rate_limit": int(api_throttle_rate_limit),
+ "burst_limit": int(api_throttle_burst_limit),
+}
+api_quota_limit = app.node.try_get_context("API_QUOTA_LIMIT") or 25
+api_quota_period = app.node.try_get_context("API_QUOTA_PERIOD") or "DAY"
+api_quota_settings = {"limit": int(api_quota_limit), "period": api_quota_period}
+
+ApigwBedrockCognitoCdkStack(
+ app,
+ "ApigwBedrockCognitoCdkStack",
+ organization_domain=organization_domain,
+ api_throttle_settings=api_throttle_settings,
+ api_quota_settings=api_quota_settings,
+)
+
+app.synth()
diff --git a/apigw-bedrock-cognito-cdk/cdk.json b/apigw-bedrock-cognito-cdk/cdk.json
new file mode 100644
index 000000000..eb79b783a
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/cdk.json
@@ -0,0 +1,62 @@
+{
+ "app": "python3 app.py",
+ "watch": {
+ "include": [
+ "**"
+ ],
+ "exclude": [
+ "README.md",
+ "cdk*.json",
+ "requirements*.txt",
+ "source.bat",
+ "**/__init__.py",
+ "**/__pycache__",
+ "tests"
+ ]
+ },
+ "context": {
+ "@aws-cdk/aws-lambda:recognizeLayerVersion": true,
+ "@aws-cdk/core:checkSecretUsage": true,
+ "@aws-cdk/core:target-partitions": [
+ "aws",
+ "aws-cn"
+ ],
+ "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true,
+ "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true,
+ "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true,
+ "@aws-cdk/aws-iam:minimizePolicies": true,
+ "@aws-cdk/core:validateSnapshotRemovalPolicy": true,
+ "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true,
+ "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true,
+ "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true,
+ "@aws-cdk/aws-apigateway:disableCloudWatchRole": true,
+ "@aws-cdk/core:enablePartitionLiterals": true,
+ "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true,
+ "@aws-cdk/aws-iam:standardizedServicePrincipals": true,
+ "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true,
+ "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true,
+ "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true,
+ "@aws-cdk/aws-route53-patters:useCertificate": true,
+ "@aws-cdk/customresources:installLatestAwsSdkDefault": false,
+ "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true,
+ "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true,
+ "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true,
+ "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true,
+ "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true,
+ "@aws-cdk/aws-redshift:columnId": true,
+ "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true,
+ "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true,
+ "@aws-cdk/aws-apigateway:requestValidatorUniqueId": true,
+ "@aws-cdk/aws-kms:aliasNameRef": true,
+ "@aws-cdk/aws-autoscaling:generateLaunchTemplateInsteadOfLaunchConfig": true,
+ "@aws-cdk/core:includePrefixInUniqueNameGeneration": true,
+ "@aws-cdk/aws-efs:denyAnonymousAccess": true,
+ "@aws-cdk/aws-opensearchservice:enableOpensearchMultiAzWithStandby": true,
+ "@aws-cdk/aws-lambda-nodejs:useLatestRuntimeVersion": true,
+ "@aws-cdk/aws-efs:mountTargetOrderInsensitiveLogicalId": true,
+ "@aws-cdk/aws-rds:auroraClusterChangeScopeOfInstanceParameterGroupWithEachParameters": true,
+ "@aws-cdk/aws-appsync:useArnForSourceApiAssociationIdentifier": true,
+ "@aws-cdk/aws-rds:preventRenderingDeprecatedCredentials": true,
+ "@aws-cdk/aws-codepipeline-actions:useNewDefaultBranchForCodeCommitSource": true
+ }
+}
diff --git a/apigw-bedrock-cognito-cdk/example-pattern.json b/apigw-bedrock-cognito-cdk/example-pattern.json
new file mode 100644
index 000000000..b68e4b0b1
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/example-pattern.json
@@ -0,0 +1,71 @@
+{
+ "title": "API Gateway with Cognito to Lambda to Amazon Bedrock",
+ "description": "Create an AWS CDK Python application for Bedrock API with Cognito user management, domain restriction, API request throttling, and quota limits.",
+ "language": "Python",
+ "level": "200",
+ "framework": "CDK",
+ "introBox": {
+ "headline": "How it works",
+ "text": [
+ "This sample project demonstrates how to create an API Gateway with AWS CDK, backed by a Cognito user pool for user management.",
+ "Users can register and log in via an endpoint. Registration can be restricted to a specific domain, such as '@example.com'. The application supports request throttling and quota limits to the API.",
+ "Configuration for domain restriction, request throttling, and quota management can be done using CDK context variables."
+ ]
+ },
+ "gitHub": {
+ "template": {
+ "repoURL": "https://github.com/aws-samples/serverless-patterns/tree/main/apigw-bedrock-cognito-cdk",
+ "templateURL": "serverless-patterns/apigw-bedrock-cognito-cdk",
+ "projectFolder": "apigw-bedrock-cognito-cdk",
+ "templateFile": "app.py"
+ }
+ },
+ "resources": {
+ "bullets": [
+ {
+ "text": "Amazon API Gateway - REST API",
+ "link": "https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-rest-api.html"
+ },
+ {
+ "text": "Amazon Bedrock",
+ "link": "https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html"
+ },
+ {
+ "text": "Amazon Bedrock - Inference parameters for foundation models",
+ "link": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html"
+ },
+ {
+ "text": "AWS CDK v2 - Python",
+ "link": "https://docs.aws.amazon.com/cdk/v2/guide/work-with-cdk-python.html"
+ },
+ {
+ "text": "Amazon Cognito - User Pools",
+ "link": "https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-identity-pools.html"
+ }
+ ]
+ },
+ "deploy": {
+ "text": [
+ "Deploy the stack: cdk deploy
",
+ "Deploy the stack with context variables: cdk deploy --context ORGANIZATION_DOMAIN=@example.com --context API_THROTTLE_RATE_LIMIT=1 --context API_THROTTLE_BURST_LIMIT=2 --context API_QUOTA_LIMIT=25 --contect API_QUOTA_LIMIT=DAY
"
+ ]
+ },
+ "testing": {
+ "text": [
+ "Run end-to-end tests: pytest
"
+ ]
+ },
+ "cleanup": {
+ "text": [
+ "Delete the stack: cdk destroy
."
+ ]
+ },
+ "authors": [
+ {
+ "name": "Inthuson Anandakumar",
+ "image": "https://avatars.githubusercontent.com/u/8601106?v=4",
+ "bio": "Inthuson Anandakumar Technical Account Manager at AWS based out of London. Inthu currently works with strategic financial services customers for AWS, helping them architect, build, and optimise cloud applications. With a background in software development and solution architecture, he thrives on crafting innovative solutions and enjoys sharing his discoveries with customers across a wide array of sectors.",
+ "linkedin": "inthuson"
+ }
+ ]
+}
diff --git a/apigw-bedrock-cognito-cdk/requirements-dev.txt b/apigw-bedrock-cognito-cdk/requirements-dev.txt
new file mode 100644
index 000000000..171a8d854
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/requirements-dev.txt
@@ -0,0 +1,5 @@
+pytest
+pytest-dependency
+requests
+PyJWT
+boto3
\ No newline at end of file
diff --git a/apigw-bedrock-cognito-cdk/requirements.txt b/apigw-bedrock-cognito-cdk/requirements.txt
new file mode 100644
index 000000000..e077888ca
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/requirements.txt
@@ -0,0 +1,2 @@
+aws-cdk-lib==2.122.0
+constructs>=10.0.0,<11.0.0
diff --git a/apigw-bedrock-cognito-cdk/source.bat b/apigw-bedrock-cognito-cdk/source.bat
new file mode 100644
index 000000000..9e1a83442
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/source.bat
@@ -0,0 +1,13 @@
+@echo off
+
+rem The sole purpose of this script is to make the command
+rem
+rem source .venv/bin/activate
+rem
+rem (which activates a Python virtualenv on Linux or Mac OS X) work on Windows.
+rem On Windows, this command just runs this batch file (the argument is ignored).
+rem
+rem Now we don't need to document a Windows command for activating a virtualenv.
+
+echo Executing .venv\Scripts\activate.bat for you
+.venv\Scripts\activate.bat
diff --git a/apigw-bedrock-cognito-cdk/src/auth/auth.py b/apigw-bedrock-cognito-cdk/src/auth/auth.py
new file mode 100644
index 000000000..31940bf21
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/src/auth/auth.py
@@ -0,0 +1,165 @@
+import os
+import boto3
+import json
+import logging
+
+
+# Configure logging
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+# Initialize AWS clients
+cognito_client = boto3.client("cognito-idp")
+apigateway_client = boto3.client("apigateway")
+ssm = boto3.client("ssm")
+
+
+# Method to get a parameter from AWS Systems Manager Parameter Store
+def get_from_parameter_store(key):
+ parameter = ssm.get_parameter(Name=key)
+ return parameter["Parameter"]["Value"]
+
+
+# Method to validate the presence of required parameters in a request body
+def validate_parameters(body, required_params):
+ return all(param in body and len(body[param]) > 1 for param in required_params)
+
+
+# Method to generate a standardised API response
+def generate_response(status_code, success, message, data=None):
+ response_data = {"success": success, "message": message}
+ if data is not None:
+ response_data["data"] = data
+ return {"statusCode": status_code, "body": json.dumps(response_data)}
+
+
+# Method to redact sensitive information in the request body
+def redact_password(body):
+ redacted_body = body.copy()
+ redacted_body["password"] = "REDACTED"
+ return redacted_body
+
+
+# Method to log a redacted version of the event
+def log_redacted_event(event):
+ redacted_body = redact_password(json.loads(event["body"]))
+ redacted_event = event.copy()
+ redacted_event["body"] = json.dumps(redacted_body)
+ logger.info("Event: %s", redacted_event)
+
+
+# Method to create an API key and associate it with a usage plan
+def create_api_key(user_id):
+ response = apigateway_client.create_api_key(
+ name=user_id,
+ enabled=True,
+ )
+ api_key_id = response["id"]
+ apigateway_client.create_usage_plan_key(
+ usagePlanId=USAGE_PLAN_ID,
+ keyId=api_key_id,
+ keyType="API_KEY",
+ )
+ logger.info(f"API Key created for user_id: {user_id}.")
+ return response["value"]
+
+
+# Method to update the custom:api_key field in Cognito user attributes
+def update_cognito_user_api_key(email, api_key):
+ cognito_client.admin_update_user_attributes(
+ UserPoolId=USER_POOL_ID,
+ Username=email,
+ UserAttributes=[
+ {
+ "Name": "custom:api_key",
+ "Value": api_key,
+ }
+ ],
+ )
+ logger.info(f"Updated custom:api_key field for email: {email}.")
+
+
+# Method to create a new Cognito user
+def create_cognito_user(email, password, fullname):
+ response = cognito_client.sign_up(
+ ClientId=USER_POOL_CLIENT_ID,
+ Username=email,
+ Password=password,
+ UserAttributes=[
+ {"Name": "name", "Value": fullname},
+ {"Name": "email", "Value": email},
+ ],
+ )
+ logger.info(f"Created user: {email}.")
+ return response
+
+
+# Method to authenticate a Cognito user
+def login_cognito_user(email, password):
+ response = cognito_client.initiate_auth(
+ AuthFlow="USER_PASSWORD_AUTH",
+ AuthParameters={
+ "USERNAME": email,
+ "PASSWORD": password,
+ },
+ ClientId=USER_POOL_CLIENT_ID,
+ )
+ logger.info(f"Authenticated: {email}")
+ return response
+
+
+CONSTRUCT_ID = os.environ["CONSTRUCT_ID"]
+# Load parameters from AWS Systems Manager Parameter Store
+USER_POOL_CLIENT_ID = get_from_parameter_store(f"/{CONSTRUCT_ID}/user_pool/client_id")
+USER_POOL_ID = get_from_parameter_store(f"/{CONSTRUCT_ID}/user_pool/id")
+USAGE_PLAN_ID = get_from_parameter_store(f"/{CONSTRUCT_ID}/api/usage_plan/id")
+
+
+def handler(event, context):
+ try:
+ # Log a redacted version of the incoming event
+ log_redacted_event(event)
+
+ # Parse the request body and path from the event
+ body = json.loads(event["body"])
+ path = event["path"]
+
+ # Validate required parameters in the request body
+ if not validate_parameters(body, ["email", "password"]) or (
+ path == "/register" and not validate_parameters(body, ["fullname"])
+ ):
+ logger.error("Missing required parameters: email or password.")
+ return generate_response(422, False, "Missing required parameters.")
+
+ else:
+ email = body["email"]
+ password = body["password"]
+ fullname = body.get("fullname", "")
+
+ # Handle different paths from the API Gateway
+ if path == "/login":
+ # Login path, authenticate the user
+ response = login_cognito_user(email, password)
+ return generate_response(200, True, response["AuthenticationResult"])
+
+ elif path == "/register":
+ # Registration path, create a new Cognito user and associated API key
+ response = create_cognito_user(email, password, fullname)
+ api_key = create_api_key(response["UserSub"])
+ update_cognito_user_api_key(email, api_key)
+ return generate_response(
+ 200, True, f"User {email} created successfully.", {"API Key": api_key}
+ )
+
+ else:
+ # Handle any other request method
+ logger.error("Invalid method: %s.", event["httpMethod"])
+ return generate_response(400, False, "Invalid method")
+
+ except Exception as err:
+ logger.exception(f"Failed to {path[1:]}.")
+ return generate_response(
+ err.response["ResponseMetadata"]["HTTPStatusCode"],
+ False,
+ err.response["Error"]["Message"],
+ )
diff --git a/apigw-bedrock-cognito-cdk/src/bedrock/bedrock.py b/apigw-bedrock-cognito-cdk/src/bedrock/bedrock.py
new file mode 100644
index 000000000..c4d6071ae
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/src/bedrock/bedrock.py
@@ -0,0 +1,112 @@
+import json
+import boto3
+import logging
+from jsonpath_ng.ext import parse
+
+
+# Configure logging
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+# Initialize AWS clients
+bedrock = boto3.client("bedrock")
+bedrock_runtime = boto3.client("bedrock-runtime")
+
+# Mapping of model types to JSONPath expressions for response extraction
+MODEL_REPLY_MAPPING = {
+ "anthropic": "$.completion",
+ "ai21": "$.completions[0].data.text",
+ "amazon": "$.results[0].outputText",
+ "meta": "$.generation",
+}
+
+
+# Method to generate a standardised API response
+def generate_response(status_code, success, message, data=[]):
+ return {
+ "statusCode": status_code,
+ "body": json.dumps(
+ {
+ "success": success,
+ "message": message,
+ "data": data,
+ }
+ ),
+ }
+
+
+# Method to list foundation models
+def list_foundation_models():
+ foundation_models = bedrock.list_foundation_models()
+ return [
+ {"modelName": model["modelName"], "modelId": model["modelId"]}
+ for model in foundation_models["modelSummaries"]
+ ]
+
+
+# Method to invoke foundation models
+def invoke_foundation_model(model_id, inference_parameters):
+ # Invoke foundation model
+ bedrock_runtime_body = json.dumps(inference_parameters)
+ response = bedrock_runtime.invoke_model(body=bedrock_runtime_body, modelId=model_id)
+ response_body = json.loads(response["body"].read())
+ logger.info("Response: %s", response_body)
+
+ # Get reply from response body from foundation model
+ if model_id.startswith(tuple(MODEL_REPLY_MAPPING.keys())):
+ json_path = MODEL_REPLY_MAPPING[model_id.split(".")[0]]
+ jsonpath_expr = parse(json_path)
+ answer = jsonpath_expr.find(response_body)[0].value
+
+ else:
+ answer = response_body
+
+ return generate_response(
+ 200, True, "Successfully retrieved response from foundation model.", answer
+ )
+
+
+def handler(event, context):
+ try:
+ # Log incoming event
+ logger.info("Event: %s", event)
+
+ if event["httpMethod"] == "GET":
+ # Handle GET request to retrieve foundation models list
+ foundation_models = list_foundation_models()
+ logger.info("Foundation Models: %s", foundation_models)
+ return generate_response(
+ 200,
+ True,
+ "Successfully retrieved foundation models list.",
+ foundation_models,
+ )
+
+ elif event["httpMethod"] == "POST":
+ # Handle POST request to invoke foundation models with inference parameters
+ body = json.loads(event["body"])
+ model_id = body.get("modelId", None)
+ inference_parameters = body.get("inferenceParameters", None)
+
+ if not (model_id and inference_parameters):
+ logger.info(
+ "Missing required parameters, model_id: %s, inferenceParameters: %s",
+ model_id,
+ inference_parameters,
+ )
+ return generate_response(422, False, "Missing required parameters.")
+
+ return invoke_foundation_model(model_id, inference_parameters)
+
+ else:
+ # Handle any other request method
+ logger.error("Invalid method: %s.", event["httpMethod"])
+ return generate_response(400, False, "Invalid method")
+
+ except Exception as err:
+ logger.exception("Failed to retrieve response from foundation model.")
+ return generate_response(
+ err.response["ResponseMetadata"]["HTTPStatusCode"],
+ False,
+ err.response["Error"]["Message"],
+ )
diff --git a/apigw-bedrock-cognito-cdk/src/bedrock/requirements.txt b/apigw-bedrock-cognito-cdk/src/bedrock/requirements.txt
new file mode 100644
index 000000000..5b0622665
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/src/bedrock/requirements.txt
@@ -0,0 +1 @@
+jsonpath-ng
\ No newline at end of file
diff --git a/apigw-bedrock-cognito-cdk/src/pre_signup/pre_signup.py b/apigw-bedrock-cognito-cdk/src/pre_signup/pre_signup.py
new file mode 100644
index 000000000..b68c6349f
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/src/pre_signup/pre_signup.py
@@ -0,0 +1,32 @@
+import os
+import logging
+
+
+# Configure logging
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+# Get organization domain from environment variable
+ORGANIZATION_DOMAIN = os.environ.get("ORGANIZATION_DOMAIN")
+
+
+def handler(event, context):
+ # Log the incoming event
+ logger.info(event)
+
+ # Extract email domain
+ email_domain = event["request"]["userAttributes"]["email"].split("@")[1]
+
+ # Check if email domain matches organization
+ if email_domain == ORGANIZATION_DOMAIN:
+ # Automatically confirm user registration and verify email
+ event["response"]["autoConfirmUser"] = True
+ event["response"]["autoVerifyEmail"] = True
+
+ else:
+ # Raise an exception if email is not part of the organization domain
+ raise Exception(
+ f"Cannot register user as email is not part of domain: {ORGANIZATION_DOMAIN}"
+ )
+
+ return event
diff --git a/apigw-bedrock-cognito-cdk/tests/__init__.py b/apigw-bedrock-cognito-cdk/tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/apigw-bedrock-cognito-cdk/tests/e2e/__init__.py b/apigw-bedrock-cognito-cdk/tests/e2e/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/apigw-bedrock-cognito-cdk/tests/e2e/test_apigw_bedrock_cognito_cdk_stack.py b/apigw-bedrock-cognito-cdk/tests/e2e/test_apigw_bedrock_cognito_cdk_stack.py
new file mode 100644
index 000000000..260925f97
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/tests/e2e/test_apigw_bedrock_cognito_cdk_stack.py
@@ -0,0 +1,359 @@
+import pytest
+import requests
+import jwt
+import boto3
+import time
+
+
+# Update the following variables from the CDK Stack Outputs
+API_ENDPOINT = ""
+USER_POOL_ID = ""
+
+# Test variables
+TEST_EMAIL = "johndoe@example.com"
+NON_ORG_TEST_EMAIL = None
+TEST_FULLNAME = "John Doe"
+TEST_PASSWORD = "HelloWorld123!"
+
+
+class TestClass(object):
+ # Helper function to setup the test class
+ def setup_class(self):
+ self.cognito_client = boto3.client("cognito-idp")
+
+ # Helper function to tear down created users
+ def teardown_class(self):
+ self.cognito_client.admin_delete_user(
+ UserPoolId=USER_POOL_ID, Username=TEST_EMAIL
+ )
+
+ # Helper function to handle API requests
+ def make_request(self, path, method="post", json=None, headers=None):
+ url = API_ENDPOINT + path
+ method = method.lower()
+
+ if method == "post":
+ response = requests.post(url, json=json, headers=headers)
+ elif method == "get":
+ response = requests.get(url, headers=headers)
+ else:
+ raise ValueError("Invalid HTTP method")
+
+ return response
+
+ # Helper function to check common assertions
+ def assert_common_response(
+ self, response, expected_status_code, expected_message, expected_success=None
+ ):
+ assert response.status_code == expected_status_code
+
+ response = response.json()
+
+ assert response["message"] == expected_message
+ if expected_success is not None:
+ assert response["success"] == expected_success
+
+ # Helper function to confirm user status in Cognito
+ @pytest.fixture(scope="session")
+ def confirm_user_status(self):
+ try:
+ self.cognito_client.admin_confirm_sign_up(
+ UserPoolId=USER_POOL_ID, Username=TEST_EMAIL
+ )
+ except self.cognito_client.exceptions.NotAuthorizedException as error:
+ if (
+ error.response["Error"]["Code"] == "NotAuthorizedException"
+ and error.response["Error"]["Message"]
+ == "User cannot be confirmed. Current status is CONFIRMED"
+ ):
+ pass
+ else:
+ raise error
+
+ # Helper function to login in and store the Id and API tokens for re-use
+ @pytest.fixture(scope="session")
+ def login(self, confirm_user_status):
+ response = self.make_request(
+ "/login", json={"email": TEST_EMAIL, "password": TEST_PASSWORD}
+ )
+
+ id_token = response.json()["message"]["IdToken"]
+ api_token = jwt.decode(id_token, options={"verify_signature": False})[
+ "custom:api_key"
+ ]
+ time.sleep(120)
+ return id_token, api_token
+
+ # Register empty user
+ def test_register_empty_user(self):
+ response = self.make_request(
+ "/register", json={"email": "", "password": "", "fullname": ""}
+ )
+ self.assert_common_response(
+ response, 422, "Missing required parameters.", False
+ )
+
+ # Register user without email
+ def test_register_user_without_email(self):
+ response = self.make_request(
+ "/register",
+ json={"email": "", "password": TEST_PASSWORD, "fullname": TEST_FULLNAME},
+ )
+ self.assert_common_response(
+ response, 422, "Missing required parameters.", False
+ )
+
+ # Register user without password
+ def test_register_user_without_password(self):
+ response = self.make_request(
+ "/register",
+ json={"email": TEST_EMAIL, "password": "", "fullname": TEST_FULLNAME},
+ )
+ self.assert_common_response(
+ response, 422, "Missing required parameters.", False
+ )
+
+ # Register user without fullname
+ def test_register_user_without_fullname(self):
+ response = self.make_request(
+ "/register",
+ json={"email": TEST_EMAIL, "password": TEST_PASSWORD, "fullname": ""},
+ )
+ self.assert_common_response(
+ response, 422, "Missing required parameters.", False
+ )
+
+ # Invalid email format
+ def test_register_invalid_email_format(self):
+ response = self.make_request(
+ "/register",
+ json={
+ "email": "helloworld",
+ "password": TEST_PASSWORD,
+ "fullname": TEST_FULLNAME,
+ },
+ )
+ self.assert_common_response(
+ response, 400, "Invalid email address format.", False
+ )
+
+ # Non org domain
+ @pytest.mark.skipif(
+ NON_ORG_TEST_EMAIL is None, reason="Non Org domain email is set to None."
+ )
+ def test_register_non_org_domain(self):
+ response = self.make_request(
+ "/register",
+ json={
+ "email": NON_ORG_TEST_EMAIL,
+ "password": TEST_PASSWORD,
+ "fullname": TEST_FULLNAME,
+ },
+ )
+ assert response.status_code == 400
+
+ response = response.json()
+
+ assert response["message"].startswith(
+ "PreSignUp failed with error Cannot register user as email is not part of domain: "
+ )
+ assert response["success"] == False
+
+ # Unregistered user login
+ def test_login_unregistered_user(self):
+ response = self.make_request(
+ "/login", json={"email": "XXXXXXXXXXXXXXXXXXXXX", "password": TEST_PASSWORD}
+ )
+ self.assert_common_response(response, 400, "User does not exist.", False)
+
+ # Register User
+ @pytest.mark.dependency(name="test_register_user")
+ def test_register_user(self):
+ response = self.make_request(
+ "/register",
+ json={
+ "email": TEST_EMAIL,
+ "password": TEST_PASSWORD,
+ "fullname": TEST_FULLNAME,
+ },
+ )
+ self.assert_common_response(
+ response, 200, f"User {TEST_EMAIL} created successfully.", True
+ )
+ assert response.json()["data"]["API Key"] is not None
+
+ # User re-register
+ def test_register_user_again(self):
+ response = self.make_request(
+ "/register",
+ json={
+ "email": TEST_EMAIL,
+ "password": TEST_PASSWORD,
+ "fullname": TEST_FULLNAME,
+ },
+ )
+ self.assert_common_response(response, 400, "User already exists", False)
+
+ # Registered user login
+ @pytest.mark.dependency(
+ name="test_login_registered_user", depends=["test_register_user"]
+ )
+ def test_login_registered_user(self, login):
+ response = self.make_request(
+ "/login", json={"email": TEST_EMAIL, "password": TEST_PASSWORD}
+ )
+ assert response.status_code == 200
+ assert response.json()["success"] is True
+ assert response.json()["message"]["AccessToken"] is not None
+ assert response.json()["message"]["RefreshToken"] is not None
+ assert response.json()["message"]["IdToken"] is not None
+ assert response.json()["message"]["ExpiresIn"] == 3600
+ assert response.json()["message"]["TokenType"] == "Bearer"
+
+ # GET on /bedrock endpoint without ID Token or API Token
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_get_gaterock(self):
+ response = self.make_request("/bedrock", method="get")
+ self.assert_common_response(response, 401, "Unauthorized")
+
+ # GET on /bedrock endpoint with API Token and without ID Token
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_get_gaterock_with_api_token(self, login):
+ _, api_token = login
+ response = self.make_request(
+ "/bedrock", method="get", headers={"x-api-key": api_token}
+ )
+ self.assert_common_response(response, 401, "Unauthorized")
+
+ # GET on /bedrock endpoint with ID Token and without API Token
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_get_gaterock_with_id_token(self, login):
+ id_token, _ = login
+ response = self.make_request(
+ "/bedrock", method="get", headers={"Authorization": id_token}
+ )
+ self.assert_common_response(response, 403, "Forbidden")
+
+ # GET on /bedrock endpoint with ID Token and API Token
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_get_gaterock_with_id_token_and_api_token(self, login):
+ id_token, api_token = login
+ response = self.make_request(
+ "/bedrock",
+ method="get",
+ headers={"x-api-key": api_token, "Authorization": id_token},
+ )
+ self.assert_common_response(
+ response, 200, "Successfully retrieved foundation models list.", True
+ )
+
+ # POST on /bedrock endpoint for AI21 Jurrasic Mid Model
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_post_gaterock_ai21_jurrasic_ultra_model(self, login):
+ id_token, api_token = login
+ response = self.make_request(
+ "/bedrock",
+ json={
+ "modelId": "ai21.j2-mid-v1",
+ "inferenceParameters": {
+ "prompt": "Hello World!",
+ "maxTokens": 200,
+ "temperature": 0.7,
+ "topP": 1,
+ },
+ },
+ headers={"x-api-key": api_token, "Authorization": id_token},
+ )
+ self.assert_common_response(
+ response,
+ 200,
+ "Successfully retrieved response from foundation model.",
+ True,
+ )
+ assert response.json()["data"] is not None
+
+ # POST on /bedrock endpoint for Anthropic Claude Instant Model
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_post_gaterock_anthropic_claude_2_1_model(self, login):
+ id_token, api_token = login
+ response = self.make_request(
+ "/bedrock",
+ json={
+ "modelId": "anthropic.claude-instant-v1",
+ "inferenceParameters": {
+ "prompt": "\n\nHuman:Hello World!\n\nAssistant:",
+ "max_tokens_to_sample": 200,
+ "stop_sequences": ["\n\nHuman:"],
+ },
+ },
+ headers={"x-api-key": api_token, "Authorization": id_token},
+ )
+ self.assert_common_response(
+ response,
+ 200,
+ "Successfully retrieved response from foundation model.",
+ True,
+ )
+ assert response.json()["data"] is not None
+
+ # POST on /bedrock endpoint for Amazon Titan Model
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_post_gaterock_amazon_titan_model(self, login):
+ id_token, api_token = login
+ response = self.make_request(
+ "/bedrock",
+ json={
+ "modelId": "amazon.titan-text-express-v1",
+ "inferenceParameters": {"inputText": "Hello World!"},
+ },
+ headers={"x-api-key": api_token, "Authorization": id_token},
+ )
+ self.assert_common_response(
+ response,
+ 200,
+ "Successfully retrieved response from foundation model.",
+ True,
+ )
+ assert response.json()["data"] is not None
+
+ # POST on /bedrock endpoint for Meta Llama Model
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_post_gaterock_meta_llama_model(self, login):
+ id_token, api_token = login
+ response = self.make_request(
+ "/bedrock",
+ json={
+ "modelId": "meta.llama2-13b-chat-v1",
+ "inferenceParameters": {
+ "prompt": "Hello World",
+ "max_gen_len": 128,
+ "temperature": 0.1,
+ "top_p": 0.9,
+ },
+ },
+ headers={"x-api-key": api_token, "Authorization": id_token},
+ )
+ self.assert_common_response(
+ response,
+ 200,
+ "Successfully retrieved response from foundation model.",
+ True,
+ )
+ assert response.json()["data"] is not None
+
+ # POST on /bedrock endpoint without modelId
+ @pytest.mark.dependency(depends=["test_login_registered_user"])
+ def test_post_gaterock_without_modelId(self, login):
+ id_token, api_token = login
+ response = self.make_request(
+ "/bedrock",
+ json={
+ "inferenceParameters": {
+ "prompt": "Hello World",
+ },
+ },
+ headers={"x-api-key": api_token, "Authorization": id_token},
+ )
+ self.assert_common_response(
+ response, 422, "Missing required parameters.", False
+ )
diff --git a/apigw-bedrock-cognito-cdk/utils/delete_all_api_keys.sh b/apigw-bedrock-cognito-cdk/utils/delete_all_api_keys.sh
new file mode 100644
index 000000000..43804187e
--- /dev/null
+++ b/apigw-bedrock-cognito-cdk/utils/delete_all_api_keys.sh
@@ -0,0 +1,4 @@
+for api_key_id in $(aws apigateway get-api-keys --query 'items[*].id' --output text); do
+ echo "Deleting API key $api_key_id"
+ aws apigateway delete-api-key --api-key $api_key_id
+done
\ No newline at end of file