🎃

Powertools for AWS (Python) を用いた Cognito ユーザープールのカスタマイズ方法

山本隼汰 | Hayata Yamamoto

2023/12/27に公開

こんにちは。hayata-yamamoto です。

今回は、Powertools for AWS (Python) を用いて、Cognito のユーザープールのイベントで動く Lambda 関数を実装する方法を共有します。ユーザー認証に Cognito を利用していると、時折追加の処理を実装したいと思うシーンが出てきます。そのような場合に、Lambda を用いて追加の処理を実装するのですが、Powertools を用いるとどのようなメリットがあるのかをご紹介します。

実際に弊社でも、Powertools を用いて各Cognitoのトリガーを実装しており、以前に比べて格段に実装がわかりやすく、かつ書きやすくなりました。「各種イベントのスキーマを覚えるのめんどくさいな」と思っている方にはぜひ参考にしてみていただけましたら幸いです。

Powertools記事

Lambda を用いたユーザープールワークフローのカスタマイズ

公式によりますと、以下のように説明されています

Lambda 関数を作成してから、ユーザーのサインアップ、確認、およびサインイン (認証) などのユーザープール操作中に Lambda トリガーを使用してこの関数を有効化できます。認証チャレンジの追加、ユーザーの移行、検証メッセージのカスタマイズを行うことができます。

要するにこの機能は、独自実装した Lambda 関数を事前定義し、Cognito のユーザープールにその関数を呼び出す設定を追加しておくと、各イベントが実行される前/後に定義した Lambda 関数が実行されるというものです。この時、Lambda 関数は 5秒以内にレスポンスを返さないといけなかったり、ステータスコードが500~599 の場合は再試行されない、など考慮点はありますが、それらを満たせばかなり便利に使える機能です。

ざっくりいうと以下のような流れになります。

Cognito APIを直接呼び出した場合は、APIに対してどのイベントが対応しているかを把握して利用する必要があります。以下のページに詳細が書かれており、そちらを参照しながら実装を加えていくイメージです。

実際どのように使うのか？

では以降では、実際にどのように Powertools を使って関数を定義していくのかを紹介します。ここでは、実装が比較的簡単なケースとして、トークンの生成前のトリガーを用いて Cognito が発行する JWTトークンに任意の値を含め、所定の値を消す対応を行う関数を紹介します。

他の利用シーン含め、Powertools のホームページにはいくつかの利用シーンが合わせて紹介されておりますので、ご興味ある方はそちらもご確認ください

さて、実装は以下のような感じになります。トークン生成前の関数が呼ばれたらいつでも、variable_name というキーで value を格納し、email の情報をトークンから削除するような実装となっています。

from aws_lambda_powertools.utilities.data_classes import event_source
from aws_lambda_powertools.utilities.data_classes.cognito_user_pool_event import PreTokenGenerationTriggerEvent

@event_source(data_class=PreTokenGenerationTriggerEvent)
def handler(event: PreTokenGenerationTriggerEvent, _: dict) -> dict:
    logger.info(f"{event.raw_event=}")

    # powertools で実装しない場合は、dict型の event を関数の引数に定義して
    # 以下のような実装を加えることで同様の処理を実現します。
    # event['response']['claimsOverrideDetails']['claimsToAddOrOverride'] = {}
    # event['response']['claimsOverrideDetails']['claimsToSuppress'] = []


    event.response.claims_override_details.claims_to_add_or_override = {"variable_name": "value"}
    event.response.claims_override_details.claims_to_suppress = ["email"]

    return event.raw_event

実装としてはかなりシンプルですが、Powertools が提供するデータクラスにイベントオブジェクトが事前にパースされることで、エディタ上での補完が効くようになったり、この関数に渡されるイベントオブジェクトを具に知らなくても、型定義を見ながらドキュメントを確認しつつ実装を加えていくことができるようになります。（初心者に優しいですよね）

シーケンス図で表現すると以下のような流れで処理される関数となります。

おわりに

今回は、Powertools for AWS (Python) を用いて Cognito に対して設定する簡単な Lambda 関数の定義をご紹介しました。弊社ではご紹介した方法を用いながら、AWSの恩恵を受けつつ、AWSに対して難しいとエンジニアが感じてしまう点をなくしていけるよう、日々改善しながら業務を進めています。

もしご興味お持ちいただけた方は、以下の採用情報もご確認の上ぜひカジュアル面談などにお越しくださいませ。

トドケールテックブログPublication

トドケールは「ウケワタシをもっと便利にする。」をミッションにオフィスにおける郵便物管理システム「トドケール」と遠隔郵便物管理代行サービス「クラウドメール室」を提供する企業です。

Lambda を用いたユーザープールワークフローのカスタマイズ

実際どのように使うのか？

おわりに

Discussion