Two-Factor Authentication example

FastAPI is an amazing Python framework for backend development. There are many different web resources that describe what this framework is capable of, including an outstanding official documentation. And it seems that whatever you might need is already built-in. But recently I needed to add a 2fa (Two-Factor Authentication) support to one of the projects. And it turned out to be an extremely easy task to do that with FastAPI.

FastAPI supports different security schemes, including basic authentication. Here is an example of how a user can be authenticated:

from enum import Enum
from random import SystemRandom
import strings

from fastapi import Depends, FastAPI, HTTPException, Response, status
from fastapi.security import (
    HTTPAuthorizationCredentials,
    HTTPBasic,
    HTTPBasicCredentials,
    HTTPBearer,
)
from pydantic import BaseModel

class ErrorCode(Enum):
    ok = 0
    otp_required = 1
    wrong_otp = 2
    wrong_credentials = 3

class User(BaseModel):
    username: str
    password: str

class Auth(BaseModel):
    tokens: Dict[str, User] = {}

class Database(BaseModel):
    users: List[User]
    auth: Auth

class Login(BaseModel):
    user: Optional[User]
    status: ErrorCode

class LoginResponse(BaseModel):
    token: Optional[str]
    status: ErrorCode

basic_security = HTTPBasic()
db = Database(
    users=[
        User(
            username="user",
            password="pass",
        ),
    ],
    auth=Auth(),
)

async def get_db() -> Database:
    return db

async def get_login(
    credentials: HTTPBasicCredentials = Depends(basic_security),
    db: Database = Depends(get_db),
) -> Login:
    res = Login(status=ErrorCode.wrong_credentials)
    for user in db.users:
        if (
            credentials.username == user.username
            and credentials.password == user.password
        ):
            res.user = user
            res.status = ErrorCode.ok

    return res


def get_login_response(
    login: Login = Depends(get_login), db: Database = Depends(get_db)
) -> LoginResponse:
    res = LoginResponse(status=login.status)
    if login.status == ErrorCode.ok:
        token = "".join(
            SystemRandom().choice(string.ascii_uppercase + string.digits)
            for _ in range(32)
        )
        db.auth.tokens.update({token: login.user})
        res.token = token

    return res


@app.post("/auth/credentials")
async def verify(
    response: LoginResponse = Depends(get_login_response),
) -> LoginResponse:
    return response

The /auth/credentials endpoint provides user authentication using the basic authentication scheme. This endpoint returns a token upon successful authentication, which is required to access all protected endpoints. While the example generates a random string as the token, it's important to note that in real-world scenarios, a more secure token such as JWT should be used.

Next, we add the /whoami endpoint, which requires a valid token for authentication using the bearer authentication scheme. When a user provides a valid token that corresponds to a record in the database (in this case, represented as an instance of the Database type), the endpoint returns user information. Note that this example uses a simplified implementation, and in real-world scenarios, passwords should never be stored in plain-text, and credentials should only be transmitted via HTTPS. Additionally, the example's use of Database is for illustration purposes only, and a more robust database system should be used in practice.

bearer_security = HTTPBearer()

async def get_current_user(
    bearer: HTTPAuthorizationCredentials = Depends(bearer_security),
    db: Database = Depends(get_db),
) -> User:
    if bearer.credentials in db.auth.tokens:
        return db.auth.tokens[bearer.credentials]

    raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)


@app.get("/whoami")
async def whoami(user: User = Depends(get_current_user)):
    return user.username

So far, everything seems to be working correctly. But what if we want to improve the security of the app. There is a very robust way of doing it: Two-Factor Authentication (or 2fa for short). This is a two-step process that requires users to provide two pieces of information: credentials (something they know) and an ephemeral code (generated by a separate device or app and is only valid for a short period of time, usually 30 seconds). There are different ways of generating ephemeral codes, which are referred to as one-time passwords (OTP). For simplicity, in the example below, a time-based one-time password (TOTP) will be used, which can be easily generated by apps like Microsoft Authenticator. To use the Authenticator app with our backend, the backend must be able to do two things:

1. Generate and share the secret. This is necessary to register an account in the Authenticator app.
2. Verify the codes generated by the Authenticator app.

pyotp is an excellent library that implements everything we need for this task. And it also comes with examples of usage. Before showing how to use this library together with FastAPI, we need to first take a look at what changes should be made to the code above so that 2fa can be used.

There should be a way for the user to enable/disable 2fa. For that, we can add SecuritySettings model to User and create an endpoint to enable or disable this setting:

class SecuritySettings(BaseModel):
    otp_configured: bool
    secret: str

class User(BaseModel):
    username: str
    password: str
    security_settings: SecuritySettings

@app.put("/auth/otp/enable")
async def otp_enable(otp: Otp, user: User = Depends(get_current_user)):
    user.security_settings.otp_configured = otp.enabled

After enabling 2FA as an option, we'll need to create another endpoint returning the shared secret required for OTP generation and verification. To implement this functionality, we'll be utilizing the pyotp libraries. When setting up an authenticator app, users typically scan a QR code that contains the necessary information. To generate the QR code, we'll use the qrcode library.

@app.get("/auth/otp/generate")
def generate_qr_code(user: User = Depends(get_current_user)):
    totp = pyotp.TOTP(user.security_settings.secret)
    qr_code = qrcode.make(
        totp.provisioning_uri(name=user.username, issuer_name="Example app")
    )
    img_byte_arr = io.BytesIO()
    qr_code.save(img_byte_arr, format="PNG")
    img_byte_arr = img_byte_arr.getvalue()
    return Response(content=img_byte_arr, media_type="image/png")

This endpoint generates a QR code that can be scanned by the Authenticator app to register an account:

Finally, we will modify the login endpoint to require an OTP when necessary:

async def is_otp_correct(otp: Optional[str], secret: str) -> bool:
    return pyotp.TOTP(secret).now() == otp

async def get_login(
    credentials: HTTPBasicCredentials = Depends(basic_security),
    otp: Optional[str] = None,
    db: Database = Depends(get_db),
) -> Login:
    res = Login(status=ErrorCode.wrong_credentials)
    for user in db.users:
        if (
            credentials.username == user.username
            and credentials.password == user.password
        ):
            if user.security_settings.otp_configured and not await is_otp_correct(
                otp, user.security_settings.secret
            ):
                res.status = ErrorCode.wrong_otp
            else:
                res.user = user
                res.status = ErrorCode.ok

    return res

Adding an extra layer of security to your app is simple with pyotp, qrcode, and FastAPI. Following the steps above, your app can have two-factor authentication up and running in no time. For the complete code, check out this link.

Build

python3 -m venv .venv
source ./venv/bin/activate
pip install poetry
poetry install