render qrcode and send backup tokens on tfa activation

Author: dfm88

env/.env

@@ -1,5 +1,7 @@
 FASTAPI_CONFIG=development
 
+FAKE_EMAIL_SENDER=[email protected]
+
 API_V1_STR=/api/v1
 PROJECT_NAME=fastapi-2fa
 
@@ -11,8 +13,13 @@ ALGORITHM=HS256
 FERNET_KEY_TFA_TOKEN=J_TYpprFmoLlVM0MNZElt8IwEkvEEhAwCmb8P_f7Fro=
 TFA_BACKUP_TOKENS_NR=5
 TFA_TOKEN_LENGTH=6
+# default tolerance = 30 sec
+# this number is multiplied for 30 sec to increas it.
+# -->MAX = 10 => 5 minutes
+TOTP_TOKEN_DURATION=2
+TOTP_ISSUER_NAME=fastapi_2fa
 
-ACCESS_TOKEN_EXPIRE_MINUTES=15
+ACCESS_TOKEN_EXPIRE_MINUTES=30
 REFRESH_TOKEN_EXPIRE_MINUTES=1440  # 24 h
 PRE_TFA_TOKEN_EXPIRE_MINUTES=5
 

fastapi_2fa/api/endpoints/api_v1/auth.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 from fastapi import APIRouter, Body, Depends, HTTPException, Response, status
+from fastapi.responses import StreamingResponse
 from fastapi.security import OAuth2PasswordRequestForm
 from jose import jwt
 from pydantic import ValidationError
@@ -12,6 +13,9 @@
                                         get_authenticated_user_pre_tfa)
 from fastapi_2fa.core import security
 from fastapi_2fa.core.config import settings
+from fastapi_2fa.core.enums import DeviceTypeEnum
+from fastapi_2fa.core.utils import send_backup_tokens
+from fastapi_2fa.crud.device import device_crud
 from fastapi_2fa.crud.users import user_crud
 from fastapi_2fa.models.users import User
 from fastapi_2fa.schemas.token_schema import TokenPayload, TokenSchema
@@ -20,20 +24,50 @@
 auth_router = APIRouter()
 
 
-@auth_router.post("/signup", summary="Create user", response_model=UserOut)
+@auth_router.post(
+    "/signup",
+    summary="Create user",
+    responses={
+        200: {
+            "content": {"image/png": {}},
+            "description": "Returns no content or a qr code "
+                           "if tfas is enabled and device_type "
+                           "is 'code_generator'",
+        }
+    },
+)
 async def signup(
-    db: Session = Depends(get_db), form_data: UserCreate = Depends()
+    db: Session = Depends(get_db), user_data: UserCreate = Depends()
 ) -> Any:
     try:
-        user = await user_crud.create(
+        # async with db.begin_nested():
+        user = await user_crud(transaction=True).create(
             db=db,
-            user=form_data,
+            user=user_data,
         )
-        return user
+        # raise Exception
+        if user_data.tfa_enabled:
+            device, qr_code = await device_crud(transaction=True).create(
+                db=db,
+                device=user_data.device,
+                user=user
+            )
+            send_backup_tokens(user=user, device=device)
+
+            if device.device_type == DeviceTypeEnum.CODE_GENERATOR:
+                return StreamingResponse(content=qr_code, media_type="image/png")
+
+        return Response(status_code=status.HTTP_200_OK)
+
     except IntegrityError:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
-            detail=f'User with email "{form_data.email}" already exists',
+            detail=f'User with email `{user_data.email}` already exists',
+        )
+    except Exception as ex:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=str(ex),
         )
 
 

fastapi_2fa/api/endpoints/api_v1/two_factor_auth.py

@@ -1,40 +1,59 @@
 from typing import Any
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, HTTPException, Response, status
+from fastapi.responses import StreamingResponse
 from sqlalchemy.orm import Session
 
 from fastapi_2fa.api.deps.db import get_db
 from fastapi_2fa.api.deps.users import get_authenticated_user
+from fastapi_2fa.core.enums import DeviceTypeEnum
+from fastapi_2fa.core.utils import send_backup_tokens
 from fastapi_2fa.crud.device import device_crud
 from fastapi_2fa.crud.users import user_crud
 from fastapi_2fa.models.users import User
 from fastapi_2fa.schemas.device_schema import DeviceCreate
-from fastapi_2fa.schemas.user_schema import UserOut, UserUpdate
+from fastapi_2fa.schemas.user_schema import UserUpdate
 
 tfa_router = APIRouter()
 
 
 @tfa_router.put(
     "/enable_tfa",
     summary="Enable two factor authentication for registered user",
-    response_model=UserOut,
+    responses={
+        200: {
+            "content": {"image/png": {}},
+            "description": "Returns no content or a qr code "
+                           "if device_type is 'code_generator'",
+        }
+    },
 )
 async def enable_tfa(
     device: DeviceCreate,
     db: Session = Depends(get_db),
     user: User = Depends(get_authenticated_user),
 ) -> Any:
-    if not user_crud.is_tfa_enabled(user):
-        async with db.begin_nested():
-            user = await user_crud.update(
-                db=db,
-                db_obj=user,
-                obj_in=UserUpdate(tfa_enabled=True)
-            )
-            await device_crud.create(
-                db=db,
-                device=device,
-                user=user
-            )
-    await db.refresh(user)
-    return user
+    if not user_crud(transaction=True).is_tfa_enabled(user):
+        user = await user_crud.update(
+            db=db,
+            db_obj=user,
+            obj_in=UserUpdate(tfa_enabled=True)
+        )
+        device, qr_code = await device_crud(transaction=True).create(
+            db=db,
+            device=device,
+            user=user
+        )
+
+        send_backup_tokens(user=user, device=device)
+
+        if device.device_type == DeviceTypeEnum.CODE_GENERATOR:
+            return StreamingResponse(content=qr_code, media_type="image/png")
+
+        return Response(status_code=status.HTTP_200_OK)
+
+    raise HTTPException(
+        status_code=status.HTTP_400_BAD_REQUEST,
+        detail='Two factor authentication already '
+               f'active for user {user.email}'
+    )

fastapi_2fa/core/config.py

@@ -3,7 +3,8 @@
 from typing import Any, Dict, Optional
 
 from dotenv import load_dotenv
-from pydantic import AnyHttpUrl, BaseSettings, PostgresDsn, validator
+from pydantic import (AnyHttpUrl, BaseSettings, EmailStr, Field, PostgresDsn,
+                      validator)
 
 load_dotenv(dotenv_path="./env/.env")
 
@@ -12,6 +13,8 @@ class BaseConfig(BaseSettings):
     API_V1_STR: str = os.environ.get("API_V1_STR")
     PROJECT_NAME: str = os.environ.get("PROJECT_NAME")
 
+    FAKE_EMAIL_SENDER: EmailStr = os.environ.get("FAKE_EMAIL_SENDER")
+
     JWT_SECRET_KEY: str = os.environ.get("JWT_SECRET_KEY")
     JWT_SECRET_KEY_REFRESH: str = os.environ.get("JWT_SECRET_KEY_REFRESH")
     PRE_TFA_SECRET_KEY: str = os.environ.get("PRE_TFA_SECRET_KEY")
@@ -21,6 +24,15 @@ class BaseConfig(BaseSettings):
 
     TFA_BACKUP_TOKENS_NR: int = os.environ.get("TFA_BACKUP_TOKENS_NR")
     TFA_TOKEN_LENGTH: int = os.environ.get("TFA_TOKEN_LENGTH")
+    TOTP_ISSUER_NAME: str = os.environ.get("TOTP_ISSUER_NAME")
+    # default tolerance = 30 sec
+    # this number is multiplied for 30 sec to increas it.
+    # -->MAX = 10 => 5 minutes
+    TOTP_TOKEN_DURATION: int = Field(
+        default=os.environ.get("TOTP_TOKEN_DURATION", 2),
+        gt=0,
+        le=10
+    )
 
     ACCESS_TOKEN_EXPIRE_MINUTES: int = os.environ.get("ACCESS_TOKEN_EXPIRE_MINUTES", 15)
     REFRESH_TOKEN_EXPIRE_MINUTES: int = os.environ.get(

fastapi_2fa/core/two_factor_auth.py

@@ -1,7 +1,9 @@
+import io
 import random
 from typing import Iterator
 
 import pyotp
+import qrcode
 from fernet import Fernet
 
 from fastapi_2fa.core.config import settings
@@ -43,3 +45,14 @@ def get_fake_otp_tokens(
             str(random.randint(0, 9)) for _ in range(nr_digits)
         )
         yield random_otp
+
+def qr_code_from_key(encoded_key: str, user_email: str):
+    decoded_key = _fernet_decode(value=encoded_key)
+    qrcode_key = pyotp.TOTP(
+        decoded_key, interval=settings.TOTP_TOKEN_DURATION
+    ).provisioning_uri(user_email, issuer_name=settings.TOTP_ISSUER_NAME)
+    img = qrcode.make(qrcode_key)
+    buffer = io.BytesIO()
+    img.save(buffer, format="PNG")
+    buffer.seek(0)
+    return buffer

fastapi_2fa/core/utils.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+
+from pydantic import EmailStr
+
+from fastapi_2fa.core.config import settings
+from fastapi_2fa.models.device import Device
+from fastapi_2fa.models.users import User
+from fastapi_2fa.tasks.tasks import send_email
+
+
+@dataclass(slots=True, frozen=True)
+class Email:
+    to_: list[EmailStr]
+    from_: EmailStr
+    text_: str
+
+    def __repr__(self) -> str:
+        return (
+            f"Email (to: {self.to_}), "
+            f"(from: {self.from_}), "
+            f"(text: {self.text_})"
+        )
+
+
+def send_backup_tokens(user: User, device: Device):
+    email_ob = Email(
+        to_=[user.email],
+        from_=settings.FAKE_EMAIL_SENDER,
+        text_=f"Backup tokens : {device.backup_tokens}"
+    )
+    send_email(email=email_ob)
+
+
+def send_totp_token(user: User, token: str):
+    email_ob = Email(
+        to_=[user.email],
+        from_=settings.FAKE_EMAIL_SENDER,
+        text_=f"Access TOTP token : {token}"
+    )
+    send_email(email=email_ob)

fastapi_2fa/crud/device.py

@@ -1,7 +1,9 @@
+from qrcode.image.svg import SvgImage
 from sqlalchemy.orm import Session
 
+from fastapi_2fa.core.enums import DeviceTypeEnum
 from fastapi_2fa.core.two_factor_auth import (
-    create_encoded_two_factor_auth_key, get_fake_otp_tokens)
+    create_encoded_two_factor_auth_key, get_fake_otp_tokens, qr_code_from_key)
 from fastapi_2fa.crud.base_crud import CrudBase
 from fastapi_2fa.models.device import BackupToken, Device
 from fastapi_2fa.models.users import User
@@ -31,5 +33,14 @@ async def create(
         if await self.handle_commit(db):
             await db.refresh(db_device)
 
+        # generate qr_code
+        qr_code = None
+        if device.device_type == DeviceTypeEnum.CODE_GENERATOR:
+            qr_code = qr_code_from_key(
+                encoded_key=encoded_key,
+                user_email=user.email
+            )
+        return db_device, qr_code
+
 
 device_crud = DeviceCrud(model=Device)

fastapi_2fa/main.py

@@ -11,6 +11,7 @@
 )
 
 
+# routers
 app.include_router(router=router, prefix=settings.API_V1_STR)
 
 

fastapi_2fa/models/device.py

@@ -21,9 +21,26 @@ class Device(Base):
         uselist=True,
     )
 
+    def __repr__(self) -> str:
+        return (
+            f"<{self.__class__.__name__}("
+            f"id={self.id}, "
+            f"device_type={self.device_type}, "
+            f"user={self.user.full_name}, "
+            f")>"
+        )
+
 
 class BackupToken(Base):
     id = Column(Integer, primary_key=True, index=True)
     device_id = Column(Integer, ForeignKey("device.id"))
-    device = relationship("Device", back_populates="backup_tokens", uselist=True)
+    device = relationship("Device", back_populates="backup_tokens", uselist=False)
     token = Column(String(length=8))
+
+    def __repr__(self) -> str:
+        return (
+            f"<{self.__class__.__name__}("
+            f"id={self.id}, "
+            f"device={self.device}, "
+            f")>"
+        )

fastapi_2fa/models/users.py

@@ -18,6 +18,8 @@ def __repr__(self) -> str:
         return (
             f"<{self.__class__.__name__}("
             f"id={self.id}, "
+            f"email={self.email}, "
+            f"tfa={self.tfa_enabled}, "
             f"full_name={self.full_name}, "
             f")>"
         )