概要

Cognitoアクセストークン検証を行う処理についてローカルのテスト環境内でpytestの実行ができるよう、motoを用いてCognitoアクセストークンの発行処理をモックとして構築しました。

環境

Python 3.11.1

requirements.txt

boto3==1.28.33
botocore==1.31.37
moto==4.2.0
pytest==7.4.0
pytest-mock==3.12.0
requests==2.28.2

motoについて

motoはAWSの各種サービスに基づいてテストを簡単に行えるよう、AWSのサービスをモックできるPythonライブラリになります。AWSのサービスをモック化することで、実際のAWS環境には接続せずにテストを行うことができます。とても便利。

モックしたリソース

Cognitoユーザープールのモック

以下のコードにより、Cognitoユーザープール、アプリケーションクライアント、ユーザー、Cognitoアクセストークンの取得をモックしています。また、テストコード上からユーザープールID、クライアントID、トークンを参照できるようにしています。

@pytest.fixture()
def cognito_userpool(cognitoidp):
    # ユーザープールの作成
    userpool = cognitoidp.create_user_pool(
        PoolName=f"misato-app-admin-userpool-test",
    )

    # クライアントの作成
    userpool_client = cognitoidp.create_user_pool_client(
        UserPoolId=userpool["UserPool"]["Id"],
        ClientName=f"misato-app-admin-userpool-client-test",
    )

    # ユーザーの作成
    user_name = "test-user"
    password = "P@ssw0rd"
    cognitoidp.admin_create_user(
        UserPoolId=userpool["UserPool"]["Id"],
        Username=user_name,
        MessageAction="SUPPRESS",
    )
    # ユーザーの永続的なパスワードを設定
    cognitoidp.admin_set_user_password(
        UserPoolId=userpool["UserPool"]["Id"],
        Username=user_name,
        Password=password,
        Permanent=True,
    )

    # Cognitoアクセストークンの取得
    token_response = cognitoidp.admin_initiate_auth(
        UserPoolId=userpool["UserPool"]["Id"],
        ClientId=userpool_client["UserPoolClient"]["ClientId"],
        AuthFlow="ADMIN_USER_PASSWORD_AUTH",
        AuthParameters={
            "USERNAME": user_name,
            "PASSWORD": password,
        },
    )

    os.environ["COGNITO_POOL_ID"] = userpool["UserPool"]["Id"]

    # テストコードから参照できるよう出力する。
    yield {
        "token": token_response["AuthenticationResult"]["AccessToken"],
        "client_id": userpool_client["UserPoolClient"]["ClientId"],
        "pool_id": userpool["UserPool"]["Id"],
    }

検証に必要なキー取得処理をモック

本来、Cognitoアクセストークンを検証するために必要な公開鍵情報は以下の場所にあります。

https://cognito-idp.<Region>.amazonaws.com/<userPoolId>/.well-known/jwks.json

今回はローカル環境内でテストを完結させるため、motoを使用して検証に必要な情報を取得処理もモックします。以下のコードが呼び出されたときにモックされた公開鍵情報が返されます。

@pytest.fixture
def fetch_public_keys(cognito_userpool):
    pool_id = cognito_userpool["pool_id"]
    region = os.environ["AWS_DEFAULT_REGION"]

    keys_url = (
        f"https://cognito-idp.{region}.amazonaws.com/{pool_id}/.well-known/jwks.json"
    )
    response = requests.get(keys_url).json()

    # ユーザー検証に利用する鍵を取得（公開鍵）
    jwks_client = jwt.PyJWK(response["keys"][0])
    signing_key = jwks_client.key
    return signing_key

まとめ

motoを用いてCognitoのリソースをモックすることで、外部要因の影響やAWS利用料金の影響を受けることなくテストを行うことができます。ただ、少しマイナーなAWSのサービスに関するmotoを利用した情報は少ないので調査は大変でした。

コード全文

今回の調査で使用したコードの全文は以下に記載しています。

本体

import os
import jwt

REGION = os.environ["AWS_DEFAULT_REGION"]
COGNITO_POOL_ID = os.environ["COGNITO_POOL_ID"]


def _get_key(token: str) -> jwt.PyJWK:
    """
    指定したトークンに対する検証用鍵を作成する
    Parameters
    ----------
        token : str
            Cognitoアクセストークン
    """
    # パブリックJSON Webキーの一覧取得
    jwks_url = f"https://cognito-idp.{REGION}.amazonaws.com/{COGNITO_POOL_ID}/.well-known/jwks.json"
    # 今回のユーザー検証に利用する鍵を取得
    jwks_client = jwt.PyJWKClient(jwks_url)
    signing_key = jwks_client.get_signing_key_from_jwt(token)
    return signing_key.key


def check_expiration(token: str) -> bool:
    """
    指定したトークンが失効していないか確認する。

    Parameters
    ----------
        token : str
            Cognitoアクセストークン

    Returns
    -------
        bool: トークンの有効性情報（True:有効、False:無効）
    """
    verify = False

    issuer = f"https://cognito-idp.{REGION}.amazonaws.com/{COGNITO_POOL_ID}"

    # 今回のユーザー検証に利用する鍵を取得
    signing_key = _get_key(token)

    try:
        token = jwt.decode(
            token,  # トークン
            key=signing_key,  # 署名を検証するための鍵
            algorithms=["RS256"],  # 署名のアルゴリズム
            issuer=issuer,  # トークンの発行者
        )

        # token_useの検証
        if token["token_use"] == "access":
            verify = True

    except Exception as exc:
        print(str(exc))
        return False

    return verify

テストコード

import boto3
from moto import mock_cognitoidp
import pytest
import jwt
import os
import importlib
import requests


@pytest.fixture(scope="session", autouse=True)
def aws_credentials():
    # ダミークレデンシャル
    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
    os.environ["AWS_SECURITY_TOKEN"] = "testing"
    os.environ["AWS_SESSION_TOKEN"] = "testing"
    os.environ["AWS_DEFAULT_REGION"] = "ap-northeast-1"


@pytest.fixture()
def cognitoidp(aws_credentials):
    with mock_cognitoidp():
        yield boto3.client("cognito-idp", "ap-northeast-1")


@pytest.fixture()
def cognito_userpool(cognitoidp):
    # ユーザープールの作成
    userpool = cognitoidp.create_user_pool(
        PoolName=f"misato-app-admin-userpool-test",
    )

    # クライアントの作成
    userpool_client = cognitoidp.create_user_pool_client(
        UserPoolId=userpool["UserPool"]["Id"],
        ClientName=f"misato-app-admin-userpool-client-test",
    )

    # ユーザーの作成
    user_name = "test-user"
    password = "P@ssw0rd"
    cognitoidp.admin_create_user(
        UserPoolId=userpool["UserPool"]["Id"],
        Username=user_name,
        MessageAction="SUPPRESS",
    )
    # ユーザーの永続的なパスワードを設定
    cognitoidp.admin_set_user_password(
        UserPoolId=userpool["UserPool"]["Id"],
        Username=user_name,
        Password=password,
        Permanent=True,
    )

    # Cognitoアクセストークンの取得
    token_response = cognitoidp.admin_initiate_auth(
        UserPoolId=userpool["UserPool"]["Id"],
        ClientId=userpool_client["UserPoolClient"]["ClientId"],
        AuthFlow="ADMIN_USER_PASSWORD_AUTH",
        AuthParameters={
            "USERNAME": user_name,
            "PASSWORD": password,
        },
    )

    os.environ["COGNITO_POOL_ID"] = userpool["UserPool"]["Id"]

    yield {
        "token": token_response["AuthenticationResult"]["AccessToken"],
        "client_id": userpool_client["UserPoolClient"]["ClientId"],
        "pool_id": userpool["UserPool"]["Id"],
    }


# Cognitoアクセストークン検証用鍵生成情報（パブリックJSON Webキーの一覧）のモック
@pytest.fixture
def fetch_public_keys(cognito_userpool):
    pool_id = cognito_userpool["pool_id"]
    region = os.environ["AWS_DEFAULT_REGION"]

    keys_url = (
        f"https://cognito-idp.{region}.amazonaws.com/{pool_id}/.well-known/jwks.json"
    )
    response = requests.get(keys_url).json()

    # ユーザー検証に利用する鍵を取得（公開鍵）
    jwks_client = jwt.PyJWK(response["keys"][0])
    signing_key = jwks_client.key
    return signing_key


def test_check_expiration(cognito_userpool, mocker, fetch_public_keys):
    common_token = importlib.import_module("token_check")
    token = cognito_userpool["token"]

    # _get_keyメソッドをモック(検証用鍵の取得リクエストをモック)
    mocker.patch(
        "token_check._get_key",
        return_value=fetch_public_keys,
    )

    ret = common_token.check_expiration(token)
    assert ret == True

参考文献