Cognitoのトークンを取得するPythonスクリプト

Cognitoを使った認証認可周りを開発しているときに、jwtを何度も取得することになると思います。その時に使っている手元スクリプトです。シンプルですが結構使っています。

• 設定は対話形式で取得可能
• .envにも対応しています
• 必要なモジュール
  • boto3
  • python-dotenv

#!/usr/bin/env python
# ------------------------------------------
# Cognitoのトークン情報を取得
# ------------------------------------------
import base64
import hashlib
import hmac
import os
from getpass import getpass

import boto3
from dotenv import load_dotenv

# .envを読み込み
load_dotenv()


# パラメータ設定
def get_parameter(key: str) -> str:
    if default := os.getenv(key):
        return input(f"{key} default({default}): ") or default
    else:
        return input(f"{key}: ")


cognito_client_id = get_parameter(key="COGNITO_CLIENT_ID")
cognito_client_secret = get_parameter(key="COGNITO_CLIENT_SECRET")
aws_region = get_parameter(key="AWS_REGION")
username = get_parameter(key="USERNAME")
password = getpass("PASSWORD: ")


# IAM権限設定は不要です
cognito_idp_client = boto3.client("cognito-idp", region_name=aws_region)

# SECRET_HASH計算
message = bytes(username + cognito_client_id, "utf-8")
key = bytes(cognito_client_secret, "utf-8")
secret_hash = base64.b64encode(
    hmac.new(key, message, digestmod=hashlib.sha256).digest()
).decode()

# Cognitoから認証情報取得
res = cognito_idp_client.initiate_auth(
    AuthFlow="USER_PASSWORD_AUTH",
    AuthParameters={
        "USERNAME": username,
        "PASSWORD": password,
        "SECRET_HASH": secret_hash,
    },
    ClientId=cognito_client_id,
)


# Tokenの表示
id_token = res["AuthenticationResult"]["IdToken"]
print("---- ID TOKEN ----")
print(id_token)

id_token = res["AuthenticationResult"]["AccessToken"]
print("---- Access TOKEN ----")
print(id_token)

スクリプト実行すると対話形式でパラメータ入力画面になります。

$ ./utils/get_token.py
COGNITO_CLIENT_ID default(xxxxx): 
COGNITO_CLIENT_SECRET default(xxxxx): 
AWS_REGION default(ap-northeast-1): 
USERNAME default(xxxxx): 
PASSWORD:

.env ファイルに設定を入れておいてもよいです。

COGNITO_CLIENT_ID=xxxxx
COGNITO_CLIENT_SECRET=xxxxx
AWS_REGION=ap-northeast-1
USERNAME=xxxxx

認証が通ればトークンが取得できます。

---- ID TOKEN ----
eyJrxxxx
---- Access TOKEN ----
eyJrxxxx

上のスクリプトは USER_PASSWORD_AUTH に対応していますが、ALLOW_USER_SRP_AUTH に対応したい場合は以下の記事に書いているので、組み合わせると対応可能です。