TOTP handling in Python

python_totp_handling.md

Here are some basic snippets to handle TOTP tokens in Python.

Please not these are not related to Flask in anyway! You can adapt it in any context and/or framework (FWIW last time I implemented something along thoses lines was using the FastAPI framework. Check it out, it's pretty great!).

Here we go:

1. Define some sort of multi-factor authentication table in your database, holding securely the shared secret for each user. For instance using sqlalchemy, sqlalchemy-utils and cryptography:

from sqlalchemy import Table, Column, ForeignKey, MetaData, Unicode
from sqlalchemy_utils import EncryptedType, UUIDType
from sqlalchemy_utils.types.encrypted.encrypted_type import FernetEngine


...

# load your database encryption key (preferably in some app settings module)
db_encryption_key = os.environ.get("DB_ENCRYPTION_KEY")

...

# the usual sqlalchemy setup
convention = {
    "ix": "ix_%(column_0_label)s",
    "uq": "uq_%(table_name)s_%(column_0_name)s",
    "ck": "ck_%(table_name)s_%(constraint_name)s",
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s",
}

metadata = MetaData(naming_convention=convention)

...

# assumes there is a table named `users` with a primary key `id` of type `UUID`
users_mfa = Table(
    "users_mfa",
    metadata,
    Column("id", UUIDType, primary_key=True, index=True),
    Column(
        "user_id",
        UUIDType,
        ForeignKey("users.id", onupdate="CASCADE", ondelete="CASCADE"),
        nullable=False,
        unique=True,
        index=True,
    ),
    Column(
        "totp_secret_key",
        EncryptedType(
            type_in=Unicode,
            key=db_encryption_key,
            engine=FernetEngine,
            padding=None,
        ),
        nullable=False,
        unique=True,
    ),
)

2. Write a few functions to generate and validate TOTP tokens. For instance, using the pyotp package:

import pyotp


def generate_secret_key() -> str:
    secret_key: str = pyotp.random_base32(length=32)
    return secret_key


def generate_token(secret_key: str, num_digits: int = 6) -> str:
    try:
        totp = pyotp.TOTP(secret_key, digits=num_digits)
        return totp.now()
    except (TypeError, binascii.Error):
        raise ValueError("Error creating code")


def validate_token(secret_key: str, token: str, expiration_mn: int = 0, num_digits: int = 6) -> str:
    try:
        totp = pyotp.TOTP(secret_key, digits=num_digits)
        valid: bool = totp.verify(token, valid_window=expiration_mn * 2)
    except (TypeError, binascii.Error):
        raise ValueError("Error validating totp")

    if not valid:
        raise ValueError("Error validating code")

    return totp