Building a FastAPI Webpage with AWS Cognito Authentication and a Simple Chatbot

fastapi aws cognito
data engineering
Publish Date:   2023-03-28

Introduction:
In this tutorial, we will build a FastAPI webpage that uses AWS Cognito for user authentication. Once the user logs in, they will see a simple chatbot page. We will walk through the code step by step and explain the main sections.

Before we start, make sure you have the following packages installed in your Python environment:

pip install fastapi
pip install httpx
pip install pyjwt[crypto]
pip install uvicorn

Notice that, the default pyjwt package won’t work, you might get erros such as The JWK Set did not contain any usable keys or AttributeError: module 'jwt' has no attribute 'JWTError'.
It has to be the package install like this:

pip install pyjwt[crypto]

Now, let’s dive into the code.

  1. Import necessary libraries and create a FastAPI instance:
import os
import json
import httpx
from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.security import OAuth2PasswordBearer
from fastapi import Request
import jwt
from jwt import PyJWKClient

app = FastAPI()
  1. Set up AWS Cognito settings:

Replace the following variables with your own AWS Cognito settings.

COGNITO_DOMAIN = os.environ.get("COGNITO_DOMAIN")
COGNITO_CLIENT_ID = os.environ.get("COGNITO_CLIENT_ID")
COGNITO_CLIENT_SECRET = os.environ.get("COGNITO_CLIENT_SECRET")
COGNITO_REDIRECT_URI = os.environ.get("COGNITO_REDIRECT_URI")
COGNITO_REGION = os.environ.get("COGNITO_REGION")
COGNITO_USER_POOL_ID =  os.environ.get("COGNITO_USER_POOL_ID")

AUTH_URL = f"https://{COGNITO_DOMAIN}/oauth2/authorize"
TOKEN_URL = f"https://{COGNITO_DOMAIN}/oauth2/token"
LOGOUT_URL = f"https://{COGNITO_DOMAIN}/logout"
  1. Get Cognito JWT secret:

This function fetches the JSON Web Key Set (JWKS) from AWS Cognito and returns the public key in PEM format.

async def get_cognito_jwt_secret() -> str:
    JWKS_URL = f"https://cognito-idp.{COGNITO_REGION}.amazonaws.com/{COGNITO_USER_POOL_ID}/.well-known/jwks.json"

    async with httpx.AsyncClient() as client:
        response = await client.get(JWKS_URL)

    if response.status_code != 200:
        raise Exception("Failed to fetch JWKS from Cognito")

    jwks = response.json()
    for key_data in jwks["keys"]:
        if key_data["alg"] == "RS256" and key_data["use"] == "sig":
            key = jwk.construct(key_data)
            return key.to_pem().decode("utf-8")

    raise Exception("Failed to find a suitable public key in JWKS")
  1. Get token and current user:

These functions are used to get the token from the request and to decode the token to get the current user’s information.

async def get_token(request: Request):
    token = request.query_params.get("token")
    
    if not token:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Token is required")
    return token

async def get_current_user(token: str = Depends(get_token)) -> str:
    JWKS_URL = f"https://cognito-idp.{COGNITO_REGION}.amazonaws.com/{COGNITO_USER_POOL_ID}/.well-known/jwks.json"
    client = PyJWKClient(JWKS_URL)

    try:
        header = jwt.get_unverified_header(token)
        key = client.get_signing_key(header["kid"])
        public_key = key.key
        payload = jwt.decode(token, public_key, algorithms=["RS256"])
        return payload["sub"]
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Token has expired")
    except jwt.JWTClaimsError:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token claims")
    except jwt.JWTError:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid token")
  1. Login endpoint:

This endpoint renders the login page with a form that redirects the user to the AWS Cognito authentication page.

@app.get("/", response_class=HTMLResponse)
async def login():
    return f"""
    <html>
        <head>
            <title>Login</title>
        </head>
        <body>
            <form action="{AUTH_URL}" method="get">
                <input type="hidden" name="response_type" value="code" />
                <input type="hidden" name="client_id" value="{COGNITO_CLIENT_ID}" />
                <input type="hidden" name="redirect_uri" value="{COGNITO_REDIRECT_URI}" />
                <input type="submit" value="Login with AWS Cognito" />
            </form>
        </body>
    </html>
    """
  1. Callback endpoint:

This endpoint is called by AWS Cognito after the user has successfully logged in. It exchanges the authorization code for an access token and redirects the user to the chatbot page.
This callback url should be set up on the AWS side when you setup the app client. For test purpose, the call back url to set at aws could be “http://localhost:8000/callback".
Once on production, it should be some secure url begins with https.

@app.get("/callback")
async def callback(code: str):
    data = {
        "grant_type": "authorization_code",
        "client_id": COGNITO_CLIENT_ID,
        "client_secret":COGNITO_CLIENT_SECRET,
        "code": code,
        "redirect_uri": COGNITO_REDIRECT_URI,
    }
    headers = {
        "Content-Type": "application/x-www-form-urlencoded"
        }
    async with httpx.AsyncClient() as client:
        response = await client.post(TOKEN_URL, data=data, headers=headers)
        
    if response.status_code != 200:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid code")
    token = response.json()
  
    return RedirectResponse(url=f"/chatbot?token={token['access_token']}")
  1. Chatbot endpoint:

This endpoint renders the chatbot page for the authenticated user.

@app.get("/chatbot", response_class=HTMLResponse)
async def chatbot(sub: str = Depends(get_current_user)):
    return f"""
    <html>
        <head>
            <title>Chatbot</title>
        </head>
        <body>
            <h1>Welcome, {sub}!</h1>
            <p>Here you can chat with the robot.</p>
        </body>
    </html>
    """
  1. Logout endpoint:

This endpoint logs the user out and redirects them to the login page.

@app.get("/logout")
async def logout(request: Request):
    token = request.query_params.get("token")
    return RedirectResponse(url=f"{LOGOUT_URL}?client_id={COGNITO_CLIENT_ID}&logout_uri={COGNITO_REDIRECT_URI}&token={token}")
  1. Run the app using uvicorn:
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)
  1. The full code can be found here
    github link
    to run the code, simpley type
    python main.py

That’s it! You now have a FastAPI webpage with AWS Cognito authentication and a simple chatbot.

Author: robot learner
Link: https://datasciencebyexample.github.io/2023/03/28/fastpai-with-aws-cognito-for-user-authentication/
Reprint policy: All articles in this blog are used except for special statements CC BY 4.0 reprint policy. If reproduced, please indicate source robot learner !
fastapi aws cognito
  TOC