44 lines
1.1 KiB
Python

import pyotp
import qrcode
from io import BytesIO
import base64
from app.core.config import settings
def generate_totp_secret() -> str:
"""Generate a new TOTP secret."""
return pyotp.random_base32()
def get_totp_uri(secret: str, email: str) -> str:
"""Generate TOTP URI for QR code."""
return pyotp.totp.TOTP(secret).provisioning_uri(
name=email, issuer_name=settings.PROJECT_NAME
)
def generate_qr_code(secret: str, email: str) -> str:
"""Generate QR code as base64 string."""
totp_uri = get_totp_uri(secret, email)
qr = qrcode.QRCode(
version=1,
error_correction=qrcode.constants.ERROR_CORRECT_L,
box_size=10,
border=4,
)
qr.add_data(totp_uri)
qr.make(fit=True)
img = qr.make_image(fill_color="black", back_color="white")
buffered = BytesIO()
img.save(buffered)
img_str = base64.b64encode(buffered.getvalue()).decode()
return f"data:image/png;base64,{img_str}"
def verify_totp(secret: str, code: str) -> bool:
"""Verify TOTP code."""
totp = pyotp.TOTP(secret)
return totp.verify(code)