feat: add ALB+Cognito authentication and security enhancements

.gitignore

@@ -5,3 +5,4 @@ __pycache__
 data
 .env
 .DS_Store
+iac/kb-*.txt

README.md

@@ -21,19 +21,21 @@ This implementation is an evolution of the [AI Chat Accelerator implementation](
 - Implements Agentic RAG
 - Easily add additional tools for the agent to use
 - See conversation history and select to see past converations
+- Optional ALB+Cognito authentication (disabled by default, can be enabled for production)
 - Built-in auto scaling architecture (see docs below)
 - End to end observability with AgentCore GenAI observability and OpenTelemetry (OTEL)
 
 ## Usage
 
-Follow the 6 step process below for deploying this solution into your AWS account.
+Follow the 7 step process below for deploying this solution into your AWS account.
 
 1. Setup/Install prerequisites
 2. Deploy cloud infrastructure
 3. Deploy application code
 4. Upload your documents to the generated S3 bucket
 5. Trigger the Bedrock Knowledge Base sync
-6. Start chatting with your documents in the app
+6. Create a Cognito user for authentication (if authentication is enabled)
+7. Start chatting with your documents in the app
 
 ### 1. Setup/Install prerequisites
 
@@ -107,6 +109,43 @@ tags = {
 EOF
 ```
 
+#### Optional: Restrict web interface access by IP
+
+By default, the web interface is accessible from any IP address. To restrict access to specific IP addresses or CIDR blocks, add the `allowed_ips` parameter to your `terraform.tfvars`:
+
+```sh
+cat << EOF >> terraform.tfvars
+allowed_ips = [
+  "203.0.113.0/24",    # Your office network
+  "198.51.100.5/32",   # Your home IP
+  "192.0.2.0/24"       # Additional network
+]
+EOF
+```
+
+Replace the example IP addresses with your actual IP addresses or CIDR blocks. You can find your current IP address by running `curl ifconfig.me`.
+
+#### Optional: Enable authentication
+
+By default, authentication is disabled for easier development and testing. To enable ALB+Cognito authentication (recommended for production), add the `enable_authentication` parameter to your `terraform.tfvars`:
+
+```sh
+cat << EOF >> terraform.tfvars
+enable_authentication = true
+EOF
+```
+
+**Important Security Requirement**: ALB authentication with Cognito requires HTTPS. You must provide an SSL/TLS certificate via the `acm_certificate_arn` parameter when enabling authentication. This ensures that sensitive authentication information (credentials, tokens) is encrypted during transit.
+
+```sh
+cat << EOF >> terraform.tfvars
+enable_authentication = true
+acm_certificate_arn = "arn:aws:acm:region:account:certificate/certificate-id"
+EOF
+```
+
+When authentication is enabled, users must log in via Cognito to access the application.
+
 Deploy using terraform.
 
 ```sh
@@ -201,12 +240,56 @@ make sync
 
 Note that this script calls the `bedrock-agent start-ingestion-job` API. This job will need to successfully complete before the agent will be able to answer questions about your documents.
 
-### 6. Start chatting with your documents in the app
+### 6. Create a Cognito user for authentication (if authentication is enabled)
+
+If you have authentication enabled, create a Cognito user:
+
+```sh
+cd iac
+./create-user.sh
+```
+
+Follow the prompts to create a user with your email address and a temporary password.
+
+### 7. Start chatting with your documents in the app
 
 ```sh
 open $(terraform output -raw endpoint)
 ```
 
+If authentication is enabled, you'll be redirected to the Cognito login page. Use the email and temporary password you created in step 6. You'll be prompted to set a new password on first login.
+
+If authentication is disabled, you can access the application directly without any login.
+
+## Security
+
+### SSL/TLS and Authentication Requirements
+
+For production deployments, SSL/TLS encryption is strongly recommended. When enabling Cognito authentication, SSL/TLS is **required** - you cannot enable authentication without providing an ACM certificate.
+
+- **SSL without authentication**: Supported - provides encrypted communication
+- **Authentication without SSL**: **Not supported** - ALB authentication requires HTTPS for security
+- **Both SSL and authentication**: Recommended for production
+
+To configure SSL/TLS, provide an ACM certificate ARN:
+```hcl
+acm_certificate_arn = "arn:aws:acm:region:account:certificate/certificate-id"
+```
+
+### IP Access Restrictions
+
+The web interface can be restricted to specific IP addresses by setting the `allowed_ips` variable in your `terraform.tfvars` file. This creates security group rules that only allow HTTP traffic from the specified IP addresses or CIDR blocks.
+
+Example:
+```hcl
+allowed_ips = [
+  "203.0.113.0/24",    # Office network
+  "198.51.100.5/32",   # Specific IP address
+]
+```
+
+If not specified, the default value allows access from anywhere (`0.0.0.0/0`).
+
 ## Scaling
 
 This architecture can be scaled using two primary levers:

agent/client.py

@@ -2,6 +2,7 @@
 import json
 import argparse
 import uuid
+import base64
 
 parser = argparse.ArgumentParser(description="Create an agent runtime")
 parser.add_argument("--agent_runtime_arn", required=True, help="Agent runtime arn")
@@ -17,9 +18,12 @@
 
 agent_core_client = boto3.client("bedrock-agentcore")
 
+# Encode user_id with base64 for AgentCore compatibility
+user_id = base64.b64encode("test-user".encode()).decode().rstrip('=')
+
 payload = json.dumps({
     "input": {
-        "user_id": "6886c5c5ced611f1af8885b941a07a61",
+        "user_id": user_id,
         "prompt": "Who are you and what can you do?"
     }
 })
@@ -30,6 +34,7 @@
 response = agent_core_client.invoke_agent_runtime(
     agentRuntimeArn=agent_runtime_arn,
     runtimeSessionId=session_id,
+    runtimeUserId=user_id,
     payload=payload,
 )
 

deploy.sh

@@ -5,6 +5,11 @@ set -e
 export APP=$1
 export ARCH=$2
 
+# Validate required variables
+[ -z "$APP" ] && { echo "Error: APP variable is empty"; exit 1; }
+[ -z "$ARCH" ] && { echo "Error: ARCH variable is empty"; exit 1; }
+[ -z "$AWS_REGION" ] && { echo "Error: AWS_REGION variable is empty"; exit 1; }
+
 export VERSION=$(cat /dev/urandom | LC_ALL=C tr -dc 'a-zA-Z0-9' | fold -w 50 | head -n 1)
 
 # login to ECR

iac/cognito.tf

@@ -0,0 +1,64 @@
+resource "aws_cognito_user_pool" "main" {
+  count = var.enable_authentication ? 1 : 0
+  name = var.name
+
+  auto_verified_attributes = ["email"]
+  admin_create_user_config {
+    allow_admin_create_user_only = true
+  }
+
+  password_policy {
+    minimum_length    = 8
+    require_lowercase = true
+    require_numbers   = true
+    require_symbols   = true
+    require_uppercase = true
+  }
+
+  schema {
+    attribute_data_type = "String"
+    name               = "email"
+    required           = true
+    mutable            = true
+  }
+
+  tags = var.tags
+}
+
+resource "aws_cognito_user_pool_client" "main" {
+  count        = var.enable_authentication ? 1 : 0
+  name         = var.name
+  user_pool_id = aws_cognito_user_pool.main[0].id
+
+  generate_secret                      = true
+  allowed_oauth_flows_user_pool_client = true
+  allowed_oauth_flows                  = ["code"]
+  allowed_oauth_scopes                 = ["email", "openid", "profile"]
+
+  callback_urls = var.acm_certificate_arn != "" ? [
+    "https://${module.alb.dns_name}/oauth2/idpresponse"
+  ] : [
+    "http://${module.alb.dns_name}/oauth2/idpresponse"
+  ]
+
+  logout_urls = var.acm_certificate_arn != "" ? [
+    "https://${module.alb.dns_name}"
+  ] : [
+    "http://${module.alb.dns_name}"
+  ]
+
+  supported_identity_providers = ["COGNITO"]
+}
+
+resource "aws_cognito_user_pool_domain" "main" {
+  count        = var.enable_authentication ? 1 : 0
+  domain       = "${var.name}-${random_string.cognito_domain[0].result}"
+  user_pool_id = aws_cognito_user_pool.main[0].id
+}
+
+resource "random_string" "cognito_domain" {
+  count   = var.enable_authentication ? 1 : 0
+  length  = 8
+  special = false
+  upper   = false
+}

iac/create-user.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# Script to create a Cognito user for the AI Agent application
+
+set -e
+
+# Get Cognito User Pool ID from Terraform output
+USER_POOL_ID=$(terraform output -raw cognito_user_pool_id)
+
+if [ -z "$USER_POOL_ID" ]; then
+    echo "Error: Could not get Cognito User Pool ID from Terraform output"
+    exit 1
+fi
+
+# Prompt for user email
+read -p "Enter email address for the new user: " EMAIL
+
+if [ -z "$EMAIL" ]; then
+    echo "Error: Email address is required"
+    exit 1
+fi
+
+# Prompt for temporary password
+read -s -p "Enter temporary password (min 8 chars, must include uppercase, lowercase, number, symbol): " TEMP_PASSWORD
+echo
+
+if [ -z "$TEMP_PASSWORD" ]; then
+    echo "Error: Password is required"
+    exit 1
+fi
+
+echo "Creating user in Cognito User Pool: $USER_POOL_ID"
+
+# Create the user
+aws cognito-idp admin-create-user \
+    --user-pool-id "$USER_POOL_ID" \
+    --username "$EMAIL" \
+    --user-attributes Name=email,Value="$EMAIL" Name=email_verified,Value=true \
+    --temporary-password "$TEMP_PASSWORD" \
+    --message-action SUPPRESS
+
+echo "User created successfully!"
+echo "Email: $EMAIL"
+echo "Temporary password: $TEMP_PASSWORD"
+echo ""
+echo "The user will need to change their password on first login."
+echo "Access the application at: $(terraform output -raw endpoint)"