[AWS]Serverless API(五)

Lambda串接Cognito實現註冊登入

在上一篇中筆記了簡易的Cognito觸發Lambda的方式，這次反過來是由Lambda來觸發實作Cognito的功能，用途和上一篇基本上差不多，不過差異在進入點不同，用這方式就不受限於libary基本上能打API的都可以串接

不過在aws有提供各語言libary的前提下，由後端再包一次是有點多此一舉，所以這個筆記算是兼給前端使用前的AWS設定

IAM-Role

首先在IAM裡新增一個政策給lambda使用cognito和cloudwatch

{
   "Version": "2012-10-17",
   "Statement": [
      {
         "Sid": "",
         "Effect": "Allow",
         "Action": [
            "cognito-idp:AdminInitiateAuth",
            "logs:CreateLogStream",
            "logs:CreateLogGroup",
            "logs:PutLogEvents"
         ],
         "Resource": "*"
      }
   ]
}

Cognito

這邊在原本的使用者集區下新增一個應用程式用戶端

Lambda

然後是新增Lambda function，讓他綁定之前建立的Layer，使用的角色記得改用上面建立的新角色

這邊的程式碼我參考了
https://medium.com/@houzier.saurav/aws-cognito-with-python-6a2867dd02c6
並調整成FastAPI版本增加refresh_token

下面提供登入和refresh_token的程式碼，USER_POOL_ID、CLIENT_ID、CLIENT_SECRET就放入Cognito的使用者集區 ID、用戶端 ID、用戶端機密

from fastapi import FastAPI
import boto3
from mangum import Mangum
import hmac
import hashlib
import base64
import jwt

app = FastAPI()

USER_POOL_ID = 'USER_POOL_ID'
CLIENT_ID = 'CLIENT_ID'
CLIENT_SECRET = 'CLIENT_SECRET'
client = boto3.client('cognito-idp')

def get_secret_hash(username):
    msg = username + CLIENT_ID 
    dig = hmac.new(str(CLIENT_SECRET).encode('utf-8'), msg = str(msg).encode('utf-8'), digestmod=hashlib.sha256).digest()
    d2 = base64.b64encode(dig).decode()
    return d2

def initiate_auth(username, password):
    secret_hash = get_secret_hash(username)
    try:
      resp = client.admin_initiate_auth(
                 UserPoolId=USER_POOL_ID,
                 ClientId=CLIENT_ID,
                 AuthFlow='ADMIN_NO_SRP_AUTH',
                 AuthParameters={
                    'USERNAME': username,
                    'SECRET_HASH': secret_hash,
                    'PASSWORD': password,
                },
                ClientMetadata={
                  'username': username,
                  'password': password,
              })
    except client.exceptions.NotAuthorizedException:
        return None, 'The username or password is incorrect'
    except client.exceptions.UserNotConfirmedException:
        return None, 'User is not confirmed'
    except Exception as e:
        return None, e.__str__()
    return resp, None

@app.post('/cognito/signin')
async def signin(payload: dict):
    for field in ['email', 'password']:
        if payload.get(field) is None:
            return {'error': True, 'success': False, 'message': f'{field} is required', 'data': None}
    email = payload['email']
    password = payload['password']
    
    resp, msg = initiate_auth(email, password)
    if msg != None:
        return {'message': msg, 
            'error': True, 'success': False, 'data': None}
    if resp.get('AuthenticationResult'):
        id_token = resp['AuthenticationResult']['IdToken']
        decoded = jwt.decode(id_token, options={"verify_signature": False})
        return {'message': 'success', 
                'error': False, 
                'success': True, 
                'data': {
                    'sub': decoded.get('sub'),
                    'id_token': resp['AuthenticationResult']['IdToken'],
                    'refresh_token': resp['AuthenticationResult']['RefreshToken'],
                    'access_token': resp['AuthenticationResult']['AccessToken'],
                    'expires_in': resp['AuthenticationResult']['ExpiresIn'],
                    'token_type': resp['AuthenticationResult']['TokenType']
                }}
    else: #this code block is relevant only when MFA is enabled
        return {'error': True, 'success': False, 'data': None, 'message': None}
    
@app.post('/cognito/refresh')
async def refresh_token(payload: dict):
    for field in ['sub', 'token']:
        if payload.get(field) is None:
            return {'error': True, 'success': False, 'message': f'{field} is required', 'data': None}
    sub = payload['sub']
    token = payload['token']
    
    try:
        resp = client.initiate_auth(
            ClientId=CLIENT_ID,
            AuthFlow='REFRESH_TOKEN_AUTH',
            AuthParameters={
                'REFRESH_TOKEN': token,
                'SECRET_HASH': get_secret_hash(sub)
            }
        )
        if resp.get('AuthenticationResult'):
            return {'message': 'success', 
                'error': False, 
                'success': True, 
                'data': {
                    'id_token': resp['AuthenticationResult']['IdToken'],
                    'access_token': resp['AuthenticationResult']['AccessToken'],
                    'expires_in': resp['AuthenticationResult']['ExpiresIn'],
                    'token_type': resp['AuthenticationResult']['TokenType']
                }}
    except Exception as e:
        return {'error': False, 'success': True, 'message': str(e), 'data': None}

這邊有要注意的是Cognito的REFRESH_TOKEN_AUTH如果和登入一樣用mail來產生SECRET_HASH會發生下面錯誤

An error occurred (NotAuthorizedException) when calling the InitiateAuth operation: 
SecretHash does not match for the client

這是因為傳入的SECRET_HASH是需要由sub產生，無法用mail產生，只有2種refresh token才有這問題不確定是bug還是什麼原因就是…
而前端在取得token後可以在要更新時自行由token中解析出sub來用，不過我在範例中登入時順便解析了IdToken取得sub內容

這邊在Layer增加了cryptography和PyJWT這2個套件，不過要注意的是由於cryptography是有分環境版本的，如果是在win/mac環境照之前的方式打包套件的話會發現在import就發生Unable to import module的錯誤

官方有給出解決方式
https://repost.aws/zh-Hant/knowledge-center/lambda-python-package-compatible
照者指令安裝cryptography到python資料夾就可以了

pip install \    
    --platform manylinux2014_x86_64 \
    --target=./python \
    --implementation cp \
    --python-version 3.9 \
    --only-binary=:all: --upgrade \
    cryptography

接下來的api gateway連線就和之前一樣不詳述了

這樣不論正向還反向串接Cognito的方式都介紹過了，其他像註冊等程式碼可以參考下面連結的範例改寫

參考資料：

https://medium.com/@houzier.saurav/aws-cognito-with-python-6a2867dd02c6