Streamlit × Cognito でカスタム OTP ログインを実装する

cdk-streamlit-cognito/
├── app.py                 # CDK エントリポイント
├── cdk.json
├── requirements.txt       # CDK用
│
├── stacks/
│   └── main_stack.py      # CDK スタック定義
│
├── lambda/                # カスタム認証 Lambda群
│   ├── define_auth/
│   │   └── index.py
│   ├── create_auth/
│   │   └── index.py
│   └── verify_auth/
│       └── index.py
│
└── streamlit/
    ├── app.py             # Streamlit アプリ
    ├── Dockerfile
    └── requirements.txt

[ユーザー] ──(1)メール+パスワード──> [Cognito] ──(2)パスワード検証
                                        │
                                        ▼
[ユーザー] <──(3)OTPメール────────── [Lambda] ──> [SES]
                                        │
[ユーザー] ──(4)OTP入力──────────────> [Cognito]
                                        │
                                        ▼
[ユーザー] <──(5)トークン発行──────── [Lambda] ──(OTP検証)

mkdir cdk-streamlit-cognito
cd cdk-streamlit-cognito

# 仮想環境作成 & アクティベート
python -m venv .venv
source .venv/bin/activate

# 依存パッケージインストール
pip install -r requirements.txt

aws-cdk-lib>=2.170.0
constructs>=10.0.0

from aws_cdk import (
    Stack,
    CfnOutput,
    aws_cognito as cognito,
    aws_lambda as lambda_,
    aws_ec2 as ec2,
    aws_ecs as ecs,
    aws_ecs_patterns as ecs_patterns,
    aws_ecr_assets as ecr_assets,
    aws_iam as iam,
    Duration,
    RemovalPolicy,
)
from constructs import Construct

class MainStack(Stack):
    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)

        # ===================
        # Lambda Functions
        # ===================
        define_auth_fn = lambda_.Function(
            self,
            "DefineAuthChallenge",
            runtime=lambda_.Runtime.PYTHON_3_12,
            handler="index.handler",
            code=lambda_.Code.from_asset("lambda/define_auth"),
            timeout=Duration.seconds(30),
        )

        create_auth_fn = lambda_.Function(
            self,
            "CreateAuthChallenge",
            runtime=lambda_.Runtime.PYTHON_3_12,
            handler="index.handler",
            code=lambda_.Code.from_asset("lambda/create_auth"),
            timeout=Duration.seconds(30),
        )

        verify_auth_fn = lambda_.Function(
            self,
            "VerifyAuthChallenge",
            runtime=lambda_.Runtime.PYTHON_3_12,
            handler="index.handler",
            code=lambda_.Code.from_asset("lambda/verify_auth"),
            timeout=Duration.seconds(30),
        )

        # ===================
        # Cognito User Pool
        # ===================
        user_pool = cognito.UserPool(
            self,
            "UserPool",
            user_pool_name="custom-otp-pool",
            self_sign_up_enabled=True,
            sign_in_aliases=cognito.SignInAliases(email=True),
            auto_verify=cognito.AutoVerifiedAttrs(email=True),
            # Lambda トリガー
            lambda_triggers=cognito.UserPoolTriggers(
                define_auth_challenge=define_auth_fn,
                create_auth_challenge=create_auth_fn,
                verify_auth_challenge_response=verify_auth_fn,
            ),
            removal_policy=RemovalPolicy.DESTROY,
        )

        # CreateAuthChallenge Lambda に AdminGetUser と SES 権限を付与
        create_auth_fn.add_to_role_policy(
            iam.PolicyStatement(
                actions=["cognito-idp:AdminGetUser"],
                resources=[
                    f"arn:aws:cognito-idp:{self.region}:{self.account}:userpool/*"
                ],
            )
        )
        create_auth_fn.add_to_role_policy(
            iam.PolicyStatement(
                actions=["ses:SendEmail"],
                resources=["*"],
            )
        )

        # User Pool Client
        user_pool_client = user_pool.add_client(
            "AppClient",
            auth_flows=cognito.AuthFlow(
                custom=True,  # カスタム認証を有効化
                user_srp=True,  # SRP認証も有効化
            ),
            generate_secret=False,
        )

        # ===================
        # VPC
        # ===================
        vpc = ec2.Vpc(
            self,
            "Vpc",
            max_azs=2,
            nat_gateways=1,
        )

        # ===================
        # ECS Cluster
        # ===================
        cluster = ecs.Cluster(
            self,
            "Cluster",
            vpc=vpc,
        )

        # ===================
        # Docker Image
        # ===================
        image_asset = ecr_assets.DockerImageAsset(
            self,
            "StreamlitImage",
            directory="./streamlit",
        )

        # ===================
        # Fargate Service with ALB
        # ===================
        fargate_service = ecs_patterns.ApplicationLoadBalancedFargateService(
            self,
            "StreamlitService",
            cluster=cluster,
            cpu=256,
            memory_limit_mib=512,
            desired_count=1,
            task_image_options=ecs_patterns.ApplicationLoadBalancedTaskImageOptions(
                image=ecs.ContainerImage.from_docker_image_asset(image_asset),
                container_port=8501,
                environment={
                    "USER_POOL_ID": user_pool.user_pool_id,
                    "CLIENT_ID": user_pool_client.user_pool_client_id,
                    "AWS_REGION": self.region,
                },
            ),
            public_load_balancer=True,
        )

        # Cognito へのアクセス権限
        fargate_service.task_definition.task_role.add_to_policy(
            iam.PolicyStatement(
                actions=[
                    "cognito-idp:InitiateAuth",
                    "cognito-idp:RespondToAuthChallenge",
                    "cognito-idp:SignUp",
                    "cognito-idp:ConfirmSignUp",
                ],
                resources=[
                    f"arn:aws:cognito-idp:{self.region}:{self.account}:userpool/*"
                ],
            )
        )

        # ヘルスチェック設定
        fargate_service.target_group.configure_health_check(
            path="/_stcore/health",
            healthy_http_codes="200",
        )

        # ===================
        # Outputs
        # ===================
        CfnOutput(
            self,
            "UserPoolId",
            value=user_pool.user_pool_id,
            description="Cognito User Pool ID",
        )

        CfnOutput(
            self,
            "UserPoolClientId",
            value=user_pool_client.user_pool_client_id,
            description="Cognito User Pool Client ID",
        )

        CfnOutput(
            self,
            "ServiceUrl",
            value=f"http://{fargate_service.load_balancer.load_balancer_dns_name}",
            description="ALB URL",
        )

import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def handler(event, context):
    """
    認証フローの次のステップを決定する
    """
    username = event.get("userName", "unknown")
    session = event["request"]["session"]
    session_length = len(session)

    logger.info(f"User: {username}, Session step: {session_length}")

    # 空セッション → PASSWORD_VERIFIER を開始
    if session_length == 0:
        challenge = "PASSWORD_VERIFIER"
        event["response"]["challengeName"] = challenge
        event["response"]["issueTokens"] = False
        event["response"]["failAuthentication"] = False
        logger.info(f"Next challenge: {challenge}")

    # PASSWORD_VERIFIER 完了後 → OTP チャレンジ
    elif (
        session_length == 1
        and session[-1]["challengeName"] == "PASSWORD_VERIFIER"
        and session[-1]["challengeResult"]
    ):
        challenge = "CUSTOM_CHALLENGE"
        event["response"]["challengeName"] = challenge
        event["response"]["issueTokens"] = False
        event["response"]["failAuthentication"] = False
        logger.info(f"Password verified. Next challenge: {challenge}")

    # OTP 正解 → トークン発行
    elif (
        session_length == 2
        and session[-1]["challengeName"] == "CUSTOM_CHALLENGE"
        and session[-1]["challengeResult"]
    ):
        event["response"]["issueTokens"] = True
        event["response"]["failAuthentication"] = False
        logger.info("OTP verified. Issuing tokens.")

    # 失敗
    else:
        event["response"]["issueTokens"] = False
        event["response"]["failAuthentication"] = True
        logger.warning(f"Authentication failed. Session: {session}")

    return event

import logging
import random
import boto3

logger = logging.getLogger()
logger.setLevel(logging.INFO)

ses_client = boto3.client("ses")
cognito_client = boto3.client("cognito-idp")

def handler(event, context):
    """
    OTP チャレンジを作成し、メールで送信する
    """
    user_pool_id = event["userPoolId"]
    username = event["userName"]

    # ユーザーのメールアドレスを取得
    user = cognito_client.admin_get_user(UserPoolId=user_pool_id, Username=username)
    email = ""
    for attr in user["UserAttributes"]:
        if attr["Name"] == "email":
            email = attr["Value"]
            break

    # 6桁のOTPを生成
    otp = str(random.randint(100000, 999999))

    logger.info(f"User: {username}, Email: {email[:3]}***")

    # メール送信
    if email:
        try:
            ses_client.send_email(
                Source=email,  # SESで検証済みのメールアドレス
                Destination={"ToAddresses": [email]},
                Message={
                    "Subject": {"Data": "ログイン認証コード", "Charset": "UTF-8"},
                    "Body": {
                        "Text": {
                            "Data": f"あなたの認証コードは {otp} です。\n\n有効期限は3分間です。",
                            "Charset": "UTF-8",
                        }
                    },
                },
            )
            logger.info("OTP email sent successfully")
        except Exception as e:
            logger.error(f"Failed to send OTP email: {type(e).__name__}: {str(e)}")

    # 公開パラメータ（クライアントに送る）
    event["response"]["publicChallengeParameters"] = {
        "email": email[:3] + "***" + email[email.index("@"):] if email else ""
    }

    # 非公開パラメータ（検証用、クライアントには送らない）
    event["response"]["privateChallengeParameters"] = {"otp": otp}

    event["response"]["challengeMetadata"] = "OTP_CHALLENGE"

    return event

import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def handler(event, context):
    """
    OTP を検証する
    """
    username = event.get("userName", "unknown")
    expected_otp = event["request"]["privateChallengeParameters"]["otp"]
    user_otp = event["request"]["challengeAnswer"]

    # OTP が一致するか確認
    if user_otp == expected_otp:
        event["response"]["answerCorrect"] = True
        logger.info(f"User: {username}, OTP verification: SUCCESS")
    else:
        event["response"]["answerCorrect"] = False
        logger.warning(f"User: {username}, OTP verification: FAILED")

    return event

streamlit>=1.40.0
boto3>=1.35.0
pycognito>=2024.5.1

import os

import boto3
import streamlit as st
from botocore.exceptions import ClientError
from pycognito import AWSSRP

# 環境変数から取得
USER_POOL_ID = os.environ.get("USER_POOL_ID", "")
CLIENT_ID = os.environ.get("CLIENT_ID", "")
REGION = os.environ.get("AWS_REGION", "ap-northeast-1")

# Cognito クライアント
cognito_client = boto3.client("cognito-idp", region_name=REGION)

def main():
    st.title("カスタム OTP 認証デモ")

    # セッション状態の初期化
    if "authenticated" not in st.session_state:
        st.session_state.authenticated = False
    if "auth_session" not in st.session_state:
        st.session_state.auth_session = None
    if "challenge" not in st.session_state:
        st.session_state.challenge = None
    if "username" not in st.session_state:
        st.session_state.username = None
    if "confirmation_required" not in st.session_state:
        st.session_state.confirmation_required = False

    # 認証済みの場合
    if st.session_state.authenticated:
        show_dashboard()
    # メール確認待ち
    elif st.session_state.confirmation_required:
        show_confirmation()
    # OTP チャレンジ中
    elif st.session_state.challenge:
        show_challenge()
    # 未認証の場合
    else:
        show_login()

def show_login():
    """ログイン画面"""
    tab1, tab2 = st.tabs(["ログイン", "新規登録"])

    with tab1:
        st.subheader("ログイン")
        email = st.text_input("メールアドレス", key="login_email")
        password = st.text_input("パスワード", type="password", key="login_password")

        if st.button("ログイン"):
            if not email or not password:
                st.error("メールアドレスとパスワードを入力してください")
                return

            try:
                # SRP 認証を使用してカスタムチャレンジを取得
                aws_srp = AWSSRP(
                    username=email,
                    password=password,
                    pool_id=USER_POOL_ID,
                    client_id=CLIENT_ID,
                    client=cognito_client,
                )

                # CUSTOM_AUTH フローで SRP 認証を実行
                auth_params = aws_srp.get_auth_params()

                # Step 1: CUSTOM_AUTH 開始 → PASSWORD_VERIFIER が返ってくる
                response = cognito_client.initiate_auth(
                    ClientId=CLIENT_ID,
                    AuthFlow="CUSTOM_AUTH",
                    AuthParameters=auth_params,
                )

                # Step 2: PASSWORD_VERIFIER に応答 → CUSTOM_CHALLENGE が返ってくる
                if response.get("ChallengeName") == "PASSWORD_VERIFIER":
                    challenge_response = aws_srp.process_challenge(
                        response["ChallengeParameters"], auth_params
                    )
                    response = cognito_client.respond_to_auth_challenge(
                        ClientId=CLIENT_ID,
                        ChallengeName="PASSWORD_VERIFIER",
                        Session=response["Session"],
                        ChallengeResponses=challenge_response,
                    )
                else:
                    st.warning(f"予期しないチャレンジ: {response.get('ChallengeName')}")
                    return

                # カスタムチャレンジが返ってきた場合
                if response.get("ChallengeName") == "CUSTOM_CHALLENGE":
                    st.session_state.challenge = response
                    st.session_state.auth_session = response["Session"]
                    st.session_state.username = email
                    st.rerun()
                # 認証成功（チャレンジなし）
                elif "AuthenticationResult" in response:
                    st.session_state.authenticated = True
                    st.session_state.tokens = response["AuthenticationResult"]
                    st.rerun()
                else:
                    st.warning(f"予期しないレスポンス: {response.get('ChallengeName')}")

            except ClientError as e:
                st.error(f"ログインエラー: {e.response['Error']['Message']}")
            except Exception as e:
                st.error(f"エラー: {type(e).__name__}: {str(e)}")

    with tab2:
        st.subheader("新規登録")
        new_email = st.text_input("メールアドレス", key="signup_email")
        new_password = st.text_input(
            "パスワード (8文字以上、大文字・小文字・数字・記号含む)", type="password", key="signup_password"
        )

        if st.button("登録"):
            if not new_email or not new_password:
                st.error("メールアドレスとパスワードを入力してください")
                return

            try:
                # ユーザー登録
                response = cognito_client.sign_up(
                    ClientId=CLIENT_ID,
                    Username=new_email,
                    Password=new_password,
                    UserAttributes=[
                        {"Name": "email", "Value": new_email},
                    ],
                )

                st.session_state.username = new_email
                st.session_state.confirmation_required = True
                st.success("登録メールを送信しました。確認コードを入力してください。")
                st.rerun()

            except ClientError as e:
                st.error(f"登録エラー: {e.response['Error']['Message']}")

def show_confirmation():
    """メール確認画面"""
    st.subheader("メール確認")
    st.write(f"**{st.session_state.username}** に確認コードを送信しました。")

    code = st.text_input("確認コード")

    if st.button("確認"):
        try:
            cognito_client.confirm_sign_up(
                ClientId=CLIENT_ID,
                Username=st.session_state.username,
                ConfirmationCode=code,
            )

            st.success("登録完了！ログインしてください。")
            st.session_state.confirmation_required = False
            st.rerun()

        except ClientError as e:
            st.error(f"確認エラー: {e.response['Error']['Message']}")

    if st.button("確認コードを再送信"):
        try:
            cognito_client.resend_confirmation_code(
                ClientId=CLIENT_ID, Username=st.session_state.username
            )
            st.success("確認コードを再送信しました")
        except ClientError as e:
            st.error(f"再送信エラー: {e.response['Error']['Message']}")

def show_challenge():
    """OTP チャレンジ画面"""
    st.subheader("認証コード入力")

    challenge = st.session_state.challenge
    masked_email = challenge.get("ChallengeParameters", {}).get("email", "")

    st.info(f"**{masked_email}** に認証コードを送信しました。")
    otp = st.text_input("6桁の認証コードを入力してください", max_chars=6)

    col1, col2 = st.columns(2)

    with col1:
        if st.button("認証"):
            try:
                response = cognito_client.respond_to_auth_challenge(
                    ClientId=CLIENT_ID,
                    ChallengeName="CUSTOM_CHALLENGE",
                    Session=st.session_state.auth_session,
                    ChallengeResponses={
                        "USERNAME": st.session_state.username,
                        "ANSWER": otp,
                    },
                )

                # 認証成功
                if "AuthenticationResult" in response:
                    st.session_state.authenticated = True
                    st.session_state.tokens = response["AuthenticationResult"]
                    st.session_state.challenge = None
                    st.rerun()
                # さらにチャレンジがある場合
                elif "ChallengeName" in response:
                    st.session_state.challenge = response
                    st.session_state.auth_session = response["Session"]
                    st.rerun()

            except ClientError as e:
                st.error(f"認証エラー: {e.response['Error']['Message']}")

    with col2:
        if st.button("キャンセル"):
            st.session_state.challenge = None
            st.session_state.auth_session = None
            st.session_state.username = None
            st.rerun()

def show_dashboard():
    """認証後のダッシュボード"""
    # ヘッダー
    col1, col2 = st.columns([3, 1])
    with col1:
        st.title("ダッシュボード")
    with col2:
        if st.button("ログアウト", type="secondary"):
            st.session_state.authenticated = False
            st.session_state.tokens = None
            st.session_state.challenge = None
            st.session_state.auth_session = None
            st.session_state.username = None
            st.rerun()

    st.divider()

    # ユーザー情報カード
    st.subheader("ユーザー情報")
    col1, col2 = st.columns(2)
    with col1:
        st.metric(label="ログイン中のユーザー", value=st.session_state.username or "不明")
    with col2:
        st.metric(label="認証方式", value="OTP (Email)")

    st.divider()

    # 認証成功メッセージ
    st.success("カスタム OTP による多要素認証が完了しました！")

    # 機能説明
    st.subheader("このデモについて")
    st.markdown("""
    このアプリは **Amazon Cognito** のカスタム認証フローを使用して、
    以下の多要素認証を実装しています：

    1. **パスワード認証** - SRP プロトコルによる安全なパスワード検証
    2. **OTP 認証** - メールで送信される6桁のワンタイムパスワード

    すべて **Python** で構築されています：
    - インフラ: AWS CDK (Python)
    - バックエンド: AWS Lambda (Python)
    - フロントエンド: Streamlit (Python)
    """)

    # トークン情報の表示（デバッグ用）
    with st.expander("トークン情報（デバッグ用）"):
        if "tokens" in st.session_state:
            st.json(
                {
                    "AccessToken": st.session_state.tokens.get("AccessToken", "")[:50]
                    + "...",
                    "IdToken": st.session_state.tokens.get("IdToken", "")[:50] + "...",
                    "ExpiresIn": st.session_state.tokens.get("ExpiresIn"),
                    "TokenType": st.session_state.tokens.get("TokenType"),
                }
            )

if __name__ == "__main__":
    main()

FROM --platform=linux/amd64 python:3.13-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

EXPOSE 8501

CMD ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0", "--server.headless=true", "--server.enableCORS=false", "--server.enableXsrfProtection=false", "--browser.gatherUsageStats=false"]

# 仮想環境アクティベート
source .venv/bin/activate

# ブートストラップ（初回のみ）
cdk bootstrap

# デプロイ
cdk deploy

cdk destroy

aws ses verify-email-identity --email-address your-email@example.com