AWS Cognito: User Pools, JWT Flows & App Integration

User Pool
├── App Client (your React app)
├── App Client (mobile app)
├── Groups: admin, clinic-staff, billing
├── Custom Attributes: clinic_id, role
└── Lambda Triggers: pre-signup, post-confirmation, pre-token-gen

# cognito.tf

resource "aws_cognito_user_pool" "main" {
  name = "portal-users-${var.env}"

  # Sign in with email
  username_attributes      = ["email"]
  auto_verified_attributes = ["email"]

  # Password policy
  password_policy {
    minimum_length                   = 8
    require_uppercase                = true
    require_lowercase                = true
    require_numbers                  = true
    require_symbols                  = false
    temporary_password_validity_days = 7
  }

  # MFA — optional for users
  mfa_configuration = "OPTIONAL"
  software_token_mfa_configuration {
    enabled = true
  }

  # Custom attributes (can't be changed after creation!)
  schema {
    name                     = "clinic_id"
    attribute_data_type      = "String"
    mutable                  = true
    required                 = false
    string_attribute_constraints {
      min_length = 0
      max_length = 50
    }
  }

  # Email verification
  verification_message_template {
    default_email_option = "CONFIRM_WITH_CODE"
    email_subject        = "Your Learnixo verification code"
    email_message        = "Your verification code is {####}"
  }

  # Account recovery
  account_recovery_setting {
    recovery_mechanism {
      name     = "verified_email"
      priority = 1
    }
  }

  tags = { Environment = var.env }
}

# App Client for React SPA (no client secret — public client)
resource "aws_cognito_user_pool_client" "portal_web" {
  name         = "portal-web-${var.env}"
  user_pool_id = aws_cognito_user_pool.main.id

  # No secret for browser apps
  generate_secret = false

  # Token validity
  access_token_validity  = 1   # hours
  id_token_validity      = 1
  refresh_token_validity = 30  # days

  token_validity_units {
    access_token  = "hours"
    id_token      = "hours"
    refresh_token = "days"
  }

  # OAuth flows
  allowed_oauth_flows                  = ["code"]
  allowed_oauth_flows_user_pool_client = true
  allowed_oauth_scopes                 = ["email", "openid", "profile"]

  callback_urls = [
    var.env == "prod" 
      ? "https://portal.clinic.com/callback"
      : "http://localhost:5173/callback"
  ]

  logout_urls = [
    var.env == "prod"
      ? "https://portal.clinic.com"
      : "http://localhost:5173"
  ]

  supported_identity_providers = ["COGNITO"]

  # Prevent token leakage
  prevent_user_existence_errors = "ENABLED"
}

# Groups for RBAC
resource "aws_cognito_user_group" "admin" {
  name         = "admin"
  user_pool_id = aws_cognito_user_pool.main.id
  description  = "Platform administrators"
}

resource "aws_cognito_user_group" "clinic_staff" {
  name         = "clinic-staff"
  user_pool_id = aws_cognito_user_pool.main.id
  description  = "Clinic staff members"
}

{
  "sub": "a1b2c3d4-e5f6-7890-abcd-ef1234567890",
  "email": "dr.smith@sunrise-eye-care.com",
  "email_verified": true,
  "cognito:username": "dr.smith@sunrise-eye-care.com",
  "cognito:groups": ["clinic-staff"],
  "custom:clinic_id": "CLN-001",
  "iss": "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_xxxxxxx",
  "aud": "your-app-client-id",
  "exp": 1713272400,
  "iat": 1713268800
}

{
  "sub": "a1b2c3d4-...",
  "cognito:groups": ["clinic-staff"],
  "scope": "email openid profile",
  "token_use": "access",
  "exp": 1713272400
}

def get_user_context(event: dict) -> dict:
    """Extract authenticated user from API Gateway JWT authorizer context."""
    claims = event["requestContext"]["authorizer"]["jwt"]["claims"]
    
    return {
        "user_id": claims["sub"],
        "email": claims["email"],
        "clinic_id": claims.get("custom:clinic_id"),
        "groups": claims.get("cognito:groups", "").split(",") if claims.get("cognito:groups") else [],
    }

def require_group(user: dict, group: str) -> None:
    if group not in user["groups"]:
        raise PermissionError(f"User not in group: {group}")

def handler(event, context):
    user = get_user_context(event)
    
    # Only admins can delete appointments
    if event["httpMethod"] == "DELETE":
        try:
            require_group(user, "admin")
        except PermissionError:
            return {"statusCode": 403, "body": '{"error": "forbidden"}'}
    
    # Users can only access their own clinic's data
    clinic_id = event["pathParameters"]["clinic_id"]
    if user["clinic_id"] != clinic_id and "admin" not in user["groups"]:
        return {"statusCode": 403, "body": '{"error": "forbidden"}'}
    
    # ... rest of handler

# Only allow sign-ups from approved clinic domains
APPROVED_DOMAINS = {"sunrise-eye.com", "valleyoptometry.com"}

def handler(event, context):
    email = event["request"]["userAttributes"]["email"]
    domain = email.split("@")[1]
    
    if domain not in APPROVED_DOMAINS:
        raise Exception("Registration not allowed for this email domain")
    
    # Auto-confirm and auto-verify — skip email verification step
    event["response"]["autoConfirmUser"] = True
    event["response"]["autoVerifyEmail"] = True
    
    return event

# After email verification, create the user's profile in DynamoDB
import boto3, os

table = boto3.resource("dynamodb").Table(os.environ["TABLE_NAME"])

def handler(event, context):
    attrs = event["request"]["userAttributes"]
    user_id = attrs["sub"]
    
    table.put_item(Item={
        "PK": f"USER#{user_id}",
        "SK": "META",
        "type": "USER",
        "email": attrs["email"],
        "clinic_id": attrs.get("custom:clinic_id"),
        "created_at": datetime.utcnow().isoformat(),
        "status": "active"
    })
    
    return event

def handler(event, context):
    user_id = event["request"]["userAttributes"]["sub"]
    
    # Fetch additional user data from DynamoDB
    user = get_user(user_id)
    
    event["response"]["claimsOverrideDetails"] = {
        "claimsToAddOrOverride": {
            "custom:subscription": user.get("subscription_tier", "free"),
            "custom:features": ",".join(user.get("enabled_features", [])),
        }
    }
    return event

# pip install python-jose[cryptography] requests

import requests
from jose import jwk, jwt
from jose.utils import base64url_decode

REGION = "us-east-1"
USER_POOL_ID = "us-east-1_xxxxxxx"
CLIENT_ID = "your-client-id"

def get_jwks() -> list:
    url = f"https://cognito-idp.{REGION}.amazonaws.com/{USER_POOL_ID}/.well-known/jwks.json"
    return requests.get(url).json()["keys"]

# Cache JWKS — don't fetch on every request
_jwks = None

def verify_token(token: str) -> dict:
    global _jwks
    if _jwks is None:
        _jwks = get_jwks()
    
    headers = jwt.get_unverified_headers(token)
    key = next(k for k in _jwks if k["kid"] == headers["kid"])
    
    claims = jwt.decode(
        token,
        key,
        algorithms=["RS256"],
        audience=CLIENT_ID,
        options={"verify_exp": True}
    )
    
    if claims.get("token_use") not in ("access", "id"):
        raise ValueError("Invalid token_use")
    
    return claims