PythonでAWS Cognito認証

急に自分内Djangoブームが到来したので、Djangoに乗り換えはじめました。で、先日Node.jsで作ったAWSのCognito認証をPythonで作り直したので晒してみます。

事前準備

PythonでAWSにアクセスするにはIAMユーザーが必要となります。ロールにAmazonCognitoPowerUserを設定しておいてください。認証だけならロールをAmazonCognitoReadOnlyとかに落としても動くと思います。

アクセスキーも必要となりますので、認証情報タブを押してアクセスキーを作っておいてください。

つづいてAWSにPythonでアクセスするためのライブラリであるboto3を入れます

pip install boto3

AWS Cognitoで認証する

.awsに設定をしていなくても、boto3を使う時にregion等の設定を行うことができます。アプリの認証用ユーザーは固定すると思いますので、.awsを使うよりboto3を使う時に設定するほうが現実的だろうと思います。

AWS_ACCESS_KEY_IDやAWS_SECRET_ACCESS_KEY等は直接プログラムに埋め込むのは危険です(特にVCSを使ってる場合は注意)。外部ファイルに記述して読み込んで使ってください。

import boto3

def cognito_auth(user, passwd):
    # 認証開始
    try:
        aws_client = boto3.client('cognito-idp',
            region_name = "ap-northeast-1",
            aws_access_key_id = "ここにAWS_ACCESS_KEY_IDを設定します",
            aws_secret_access_key = "ここにAWS_SECRET_ACCESS_KEYを設定します",
        )

        aws_result = aws_client.admin_initiate_auth(
            UserPoolId = "ここにユーザープールIDを設定します",
            ClientId = "ここにアプリクライアントIDを設定します",
            AuthFlow = "ADMIN_NO_SRP_AUTH",
            AuthParameters = {
                "USERNAME": user,
                "PASSWORD": passwd,
            }
        )

        # 認証完了
        return aws_result

    except :
        # 認証失敗
        return None

認証が成功すれば

aws_result["AuthenticationResult"]["AccessToken"]
aws_result["AuthenticationResult"]["RefreshToken"]
aws_result["AuthenticationResult"]["IdToken"]

のようにしてトークンを拾えます。

トークンが有効かチェック

トークンはjwt形式なのでjwtライブラリで分解して……ってやろうとしたら、VC Build Tool入れて、pycrypto入れて、pyjwt入れて、winrandomの参照先変えて、decode命令呼んだら、RS256デコードしようとしてエラー吐いて死んだ／(^o^)＼

どうやってもWindowsだと暗号解除部分でエラー出すので、自力で確認するの諦めて、AWSにおまかせする方向にしました。

というわけで、保有しているトークンを使ってAWSからユーザー情報を取得します。トークンが無効ならAWSがエラーを出してくれるはずです。期限切れの確認はやってないので、もしかしたら間違ってるかもしれません。

認証成功時に

# 認証成功
req.session["aws_cognito"] = {
    "user": "認証したユーザーの名前",
    "access_token": aws_result["AuthenticationResult"]["AccessToken"],
    "refresh_token": aws_result["AuthenticationResult"]["RefreshToken"]
}

のようにして、req.sessionにトークンを保存した前提でのプログラムとなります。

import boto3

def cognito_auth_check(req):
    # 認証開始
    # *********************************************************************
    try:
        # AWSクライアント認証
        aws_client = boto3.client('cognito-idp',
            region_name = "ap-northeast-1",
            aws_access_key_id = "ここにAWS_ACCESS_KEY_IDを設定します",
            aws_secret_access_key = "ここにAWS_SECRET_ACCESS_KEYを設定します",
        )

        try:
            # 認証済みトークンの有効性確認
            return aws_client.get_user(AccessToken=req.session["aws_cognito"]["access_token"])

        except:
            # 失敗。期限切れかもしれないのでREFRESH_TOKENで再認証
            try:
                aws_result = aws_client.admin_initiate_auth(
                    UserPoolId = "ここにユーザープールIDを設定します",
                    ClientId = "ここにアプリクライアントIDを設定します",
                    AuthFlow = "REFRESH_TOKEN_AUTH",
                    AuthParameters = {
                        "USERNAME": req.session["aws_cognito"]["user"],
                        "REFRESH_TOKEN": req.session["aws_cognito"]["refresh_token"]
                    }
                )
                # アクセストークン更新
                req.session["aws_cognito"]["access_token"] = aws_result["AuthenticationResult"]["AccessToken"]

                # 再度認証確認してみる
                return aws_client.get_user(AccessToken=req.session["aws_cognito"]["access_token"])

            except:
                # REFRESH_TOKENすら期限切れや、接続障害によって失敗する可能性
                return None

    except:
        # AWSへの接続失敗。IAM変えた？　AWS接続障害の可能性も
        return None

成功すればユーザー情報が取得できます。

Djangoでデコレーターとして使う

デコレーターを作っておくとDjangoのlogin_requiredデコレーターのように

@login_required
def api(req):
    return HttpResponse('認証済みページの表示')

と、関数の前にポンッとデコレーターを置いたらページを認証が必要な状態にできるのが便利そうです。なのでデコレーターも作ってみました。

Djangoでのプログラムです。

from django.http import HttpResponse

# AWS Cognitoで認証済みであることを要求するデコレーター
def cognito_required(func):
    def wrap(req, *args, **kwargs):
        # 認証確認
        req.user = cognito_auth_check(req)
        if req.user != None:
            # デコレータ下の関数を実行
            return func(req, *args, **kwargs)

        # 失敗終了
        return HttpResponse('認証が必要です', status=401)

    wrap.__doc__=func.__doc__
    wrap.__name__=func.__name__
    return wrap

これを、account/aws_auth.pyとかで作っておけば、

from account.aws_auth import cognito_required

@cognito_required
def api(req):
    return HttpResponse('認証済みページの表示')

と、簡単に認証が必要なapiにすることができてステキ♡