Amazon Cognitoでログイン履歴を取得・管理する方法を考えてみた

import boto3
import os
import base64
import json
from typing import List, Dict, Any


def get_user_attributes_by_sub(
    user_pool_id: str, sub: str, region: str = "ap-northeast-1"
) -> List[Dict[str, Any]]:
    """
    指定したユーザプールIDとsubからユーザ属性を取得する関数
    """
    client = boto3.client("cognito-idp", region_name=region)
    try:
        response = client.admin_get_user(UserPoolId=user_pool_id, Username=sub)
        return response.get("UserAttributes", [])
    except Exception as e:
        print(f"ユーザ属性の取得中にエラー: {e}")
        raise


def lambda_handler(event: dict, context) -> dict:
    """
    Firehoseから呼び出されるLambdaのメイン関数
    """
    # リージョンを環境変数から取得（なければ東京リージョン）
    region = os.environ.get("AWS_REGION")
    if region is None:
        region = "ap-northeast-1"

    # Firehoseから渡されるレコード配列を取得
    records = event.get("records", [])
    output = []

    for record in records:
        record_id = record["recordId"]
        result = "ProcessingFailed"
        new_data = record["data"]
        try:
            # レコードのdataはBase64エンコードされているのでデコード
            encoded_data = record["data"]
            decoded_data = base64.b64decode(encoded_data).decode("utf-8")
            data_json = json.loads(decoded_data)

            # messageフィールドを取得
            message = data_json.get("message", {})

            # userSubを取得
            user_sub = message.get("userSub")

            # userPoolIdを取得（messageまたはlogSourceIdから）
            user_pool_id = message.get("userPoolId")
            if user_pool_id is None:
                log_source = data_json.get("logSourceId", {})
                user_pool_id = log_source.get("userPoolId")

            # 必須情報がなければ次のレコードへ
            if user_pool_id is None or user_sub is None:
                output.append({
                    "recordId": record_id,
                    "result": result,
                    "data": new_data
                })
                continue

            # Cognitoからemail属性を取得
            attributes = get_user_attributes_by_sub(user_pool_id, user_sub, region)
            email = None
            for attr in attributes:
                if attr.get("Name") == "email":
                    email = attr.get("Value")
                    break

            # emailをmessageに追加
            message["email"] = email
            data_json["message"] = message

            # 変換後のデータをエンコード
            new_payload = json.dumps(data_json) + "\n"
            new_data = base64.b64encode(new_payload.encode("utf-8")).decode("utf-8")
            result = "Ok"

        except Exception as e:
            # 例外発生時は失敗扱い
            print(f"レコード処理中にエラー: {e}")

        # Firehose仕様の返却形式で出力
        output.append({
            "recordId": record_id,
            "result": result,
            "data": new_data
        })

    return {"records": output}

import json
import time


def lambda_handler(event, context):
    """
    Cognito Post Authenticationトリガー用Lambdaサンプル。
    サインイン時にuserSub, email, eventType, timestampをCloudWatch Logsに出力する。
    """
    user_sub = event.get("userName")
    user_attributes = event.get("request", {}).get("userAttributes", {})
    email = user_attributes.get("email")
    event_type = "SignIn"
    timestamp = str(int(time.time()))

    log_data = {
        "userSub": user_sub,
        "email": email,
        "eventType": event_type,
        "timestamp": timestamp,
    }
    print(json.dumps(log_data))  # CloudWatch Logsに出力
    return event