【Python】CognitoのJWT検証 - 日々のあれこれφ(．．)

Cognito IDトークンの検証を実装する機会があり、やることは書いてあるけれど一通り説明できるくらいには理解しておきたかったのでその忘備録です。

gntrm.medium.com

とりわけ実装はこちらの方の記事を参考にさせていただいた形になります。

JWTとは

JSON Web Tokenとは、2者間で転送されるクレームを表すコンパクトでURLセーフな手段です。JWTのクレームは、JSON Web署名（JWS）構造のペイロードとして使用されるJSONオブジェクト、またはJSON Web暗号化（JWE）構造のプレーンテキストとしてエンコードされ、メッセージ認証コード（MAC）で暗号化してクレームをデジタル署名または整合性保護を可能にしたものになります。

JWTは以下のような構造になっています

ヘッダー：key id、tokenを暗号化に使うアルゴリズムなどが表されている。
ペイロード：ユーザー名やemailアドレスなどユーザーの情報や、認証方法、有効期限などクレームが入っている。
署名：署名は1と2をそれぞれbase64urlでエンコードして.で繋いだ文字列を、ヘッダーで指定したアルゴリズムで暗号化したもの。

# ヘッダー例
{
  "kid" : "1234example="
  "alg" : "RS256",
}

# ペイロード例
{
  "sub": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee",
  "aud": "xxxxxxxxxxxxexample",
  "iss": "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example",
  "email": "janedoe@example.com",
  "exp": 1500013000,
  "token_use": "id",
  "email_verified": true,
}

# 署名
RS256(
 base64urlEncoding(header) + '.' +
 base64urlEncoding(payload)
)

このヘッダー、ペイロード、署名をbase64urlでエンコードして . で繋いだ文字列がtokenになります

11111.22222.33333

検証の流れ

大まかな流れとしては以下です。

公開鍵を取得
署名の検証
クレームの中身の検証

JSON Web トークンの検証 - Amazon Cognito

1. 公開鍵の取得

Cognitoでは、ユーザープールごとにパブリック JSON Web キー (JWK)が2組生成されます。クライアントのtokenはそのどちらかを使って復号する形になるのですが、クライアント側のkidとパブリックkidを比較し、一致する方がそれに当たります。

from jose import jwt, jwk

# パブリックJWKを取得
def get_jwks() -> JWKS:
    return requests.get(
        f"https://cognito-idp.{os.environ.get('REGION')}.amazonaws.com/"
        f"{os.environ.get('COGNITO_POOL_ID')}/.well-known/jwks.json"
    ).json()


# kid一致を判定
def get_hmac_key(token: str, jwks: JWKS) -> Optional[JWK]:
    kid = jwt.get_unverified_header(token).get("kid")  # 1
    for key in jwks.get("keys", []):
        if key.get("kid") == kid:
            return key


token = os.environ.get('TOKEN')
hmac_key = jwk.construct(get_hmac_key(token, get_jwks()))

パブリックJWKのサンプルは以下です。

{
    "keys": [{
        "kid": "1234example=",
        "alg": "RS256",
        "kty": "RSA",
        "e": "AQAB",
        "n": "1234567890",
        "use": "sig"
    }, {
        "kid": "5678example=",
        "alg": "RS256",
        "kty": "RSA",
        "e": "AQAB",
        "n": "987654321",
        "use": "sig"
    }]
}

今回は python-jose というライブラリを使って次のようなことをしています。

kid = jwt.get_unverified_header(token).get("kid")

今回のcognitoのケースように、token発行者が複数のキーを使用していてヘッダーの情報でそれを識別しないといけないケースでは、検証の前に先にデコードします。

この処理でクライアントのkidを取得し、複数あるパブリックJWKのkidと照らし合わせて一致する方のパブリックJWKを取得します。

hmac_key = jwk.construct(get_hmac_key(token, get_jwks()))

jwk.construct(key)で与えられたJWKからキーを生成しています。

2. 署名の検証

鍵が手に入ったので、暗号化されている署名を複合化し、中身を確認します。

header_payload, encoded_signature = token.rsplit(".", 1)
decoded_signature = base64url_decode(encoded_signature.encode())
hmac_key.verify(header_payload.encode(), decoded_signature)

やっていることは単純に、署名を復号したものがヘッダー + ペイロードが一致しているか比較しています。署名は元々はヘッダー+ペイロードでできているので、パブリックjwkで復号したものが一致していれば、tokenが本物であることが証明されます。

3. クレームの中身の検証

署名を検証することで本人確認が行えましたが、本人であるからといってその本人がこのアプリケーションでリクエストを取得できる条件を満たしているとは限りません。

とりわけCognitoの場合、以下のようなものを確認する必要があります。

トークンの有効期限（exp）が切れていないこと
Audience (aud) クレームは、Amazon Cognito ユーザープールで作成されたアプリクライアント ID と一致していること
Issuer (iss) クレームは、ユーザープールと一致すること
token_use クレームが意図したものになっているか

ペイロードの中身がこれらの項目が条件を満たしているのかそれぞれを比較します。

token = os.environ.get('TOKEN')
try:
    claim = get_claim(token)
except Exception:
    raise Exception("You are not verified!")

#  有効期限が現時刻より先であることを確認
if time.time() > float(claim['exp']):
    raise Exception("token is expired!")

token検証のサンプルコード全文は以下です。