startup checks and moved tfa verify to dependency

Author: dfm88

docs/two_factor_flowchart.drawio

@@ -0,0 +1 @@
+<mxfile host="Electron" modified="2022-11-26T15:56:47.173Z" agent="5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) draw.io/20.3.0 Chrome/104.0.5112.114 Electron/20.1.3 Safari/537.36" etag="Ya5MVrSHKQGV9EB_zEQE" version="20.3.0" type="device" pages="2"><diagram id="C5RBs43oDa-KdzZeNtuy" name="Page-1">7V1bW6M4GP41ffZKH87QSw/tOIdVV+s6zk2fCGnLSkkHUrXz6zeBAIXEFnVocIoX2qThlPc7vl+CPf1k/vwpAovZ38iDQU9TvOeeftrTNFV1TPKH9qzSHls10o5p5HtsUNFx7f+CrFNhvUvfg3FpIEYowP6i3OmiMIQuLvWBKEJP5WETFJSvugBTyHVcuyDge299D8/SXkezi/4z6E9n2ZVVq59+MwfZYPYk8Qx46GmtSx/09JMIIZx+mj+fwIBOXjYvt59Xt8G3B+vTl3/in+Dm+Ovo/N+D9GTD1xySP0IEQ/zmU/9YTqK7h5vxfDJ8MlxfHU/g8IDNwiMIlmy+2LPiVTaB0CPzyZoowjM0RSEIBkXvcYSWoQfpZRTSKsZ8Q2hBOlXS+R/EeMWEAywxIl0zPA/YtzD0jijUpOkGII59N+0c+kE2JL0tei8VWLfMCRsXo2Xkwg3jdCaaIJpCvGHC1Bx4ojEQzSGOVuS4CAYA+4/lmwNMdKf5uAIe8oEh9ApB0AVoDcl0zsif2J+GywWHXoENncWnmY/h9QIkU/FEFL6MwwSFmIGkUiinFAwGa4wj9JCrEB2d6wP9ekKgOkEBipLL6h6AzsTND1v7xnIdeD/JLsZu9JVoPsIIw+e1ueYBYd8aTHdXZaP0VFgCNVPv2ZoVyA57D4JC8REhuHt9+43K9LKS1FAmbUfKtOkm16BwIwgwJH3LGEb08An1e4h4L00JiZUnU/bsx1iajr0g/HWBq60khmwt6XPQjK5uBn++Z3pZSWookyNTmTQ+kBgefbvuMNtiAA2poHGY+dTUjYZHdHJCcB8QGDhjN0Pz+2W83dCVwwNi9YZg7gfUxpzB4BFi3wUCcwgCEsdQ5AgQxAgLbSK5pB9OScsqWqNEOA6MBs2kWTaTwmBCEZhJpykzqVqtiCZaoG9OTX1Tpeqbw8FFZgu5fhp0ePDRd+kHEHr0mcEjTNTQjVYLTBUxC0se4OqPC0E03Swrl2byyqULlMtsTLl467inypXZue3atavkeONtlrNjPAHkd+rMxknjIyfIDShe2zJklU+R9zj4V+um0lKjf5XPpbMUmgSUVkBu/PieNK0p/UQJKzI7JP5L3Vri76qDZqAIRT0/7mLRqtqW/aUjUFtVoLZ5gPr7ZWAfqeQNQWYdf2lKVVqDAywnwFgICpM/98B9oMyygonwh/EfF3saWiWxkx97dgQYpyZ19EkuBcY7wX2mwGqjtlZLlWEFTQ61hATLs3FEiwB4taCfSXDyc0nLvccu8uB4CkMYAYyior8LUXKzaqk1+DJRjNIcX8YTMJqijC++JtCSiVB+0piVYsvh+Eq3VkLz9ZiUDUA1dTSh4xmi1NHR7nXLahJTRTssB566LkoYTR7VxhJGjc/7M1Q7DGthqMrHkI9FOww3YKjbrVNDveNLM2GuW43Q+lLjVd4bFquJAjT1Qz6Y2W+ulNM6w5bNlupd0sFrVA3V06UmHTrv7vY4v68Pmi7VXvKkjBtBjzLZIKApI5kPmi8+8GZzb3NArZIDGqbAYO50zYTOZwuGoo5vzo9uRmcXV59/DE67mFOIpdqvuj/TEhCluw06+RUw++z+9Lqr2HWplQedjzz32P3VB82WChpfl+9WC25MFlrn+niqOyNZqtX3o5OTwfU1GfvldkQxvvg6OOdHXQ2GV4Prs+qwznuKxMF0Du067jMn50oO1GxIJoxutQ1nYWuYYkOTaYoNXpHlRT3PPv5ODydim7bu1r45fWZnThqrrBGSWfi+3rgrzkCbxWFJqzjuA0mI1LWQOr8Yh69qVs05V+SEB3PgB11xU7QXwGibdzd4JlBTtDH15JcjmtRW4b68GhykwVvnvbfWzPpluFvhunkScU/dtqF9CKOc3eYaYGyF3ehidNlj6+voAB/Q+UqNr6QSTFOaZCvWoVapogiUqS8wnY0ttzO6fVTZTNQNb0yp9RNDtNaYbfVIKpfdTo8a1UtLtGh8p9VLgycC95m+zQzcdu2T68e61clvAk3q3lOTr3slGWGeBVTTg1JS0FU1N1U1LUt28mfyyZ+h6OPhxdXx59PTLq2rXdK0bdklTZPPEtrsEwsur6Dv7nrr7J2Yy2vcLNfdhGW9l8plh14iP8S99U1Fh4qj5z+aJZaz7IzpfbKTVEQov6t3SNWH4vnfShC3R6jeu1JJLFRV1tGu7uZsWop4yimtAZf5iy5Y2FAHlh8sWLyH6YKFN1Vw5UcLFm/Xu5r+riSiao4tOw8nJdYFrK4ukM1E3S2oltSNwxZfzs/ZzAi6iAhux2fW4DMd6Xym1S1H5dWqjv5J3Qhl8dWEFudGrQHNfm/i/D7Q+NIBe+VQ8uagalgTk8v21l5lIlj9wl5uogRwInjF797mL1Wy05Gfv/D1B0MxxucXI5LD3Jyf0mcOIIgp4C6ZaODSFDpexRjO6f14c8HWxi6iFROifekpjs1XLtrsVNtKiNpaXbveDHe1hRDt75gQtfmKSYudflsJ0fpC9d7iZz1CtL9jQjR7/iohmscZCTEqDDXoqekvGrX8JYhYtr51bW8jkiqj6pgCFyV6sVC/MRfFs3Ado/omRrUF4cbLr6LpGNVdM6pOw4wqZdjy/9+VOoXiv6Dpg/8B</diagram><diagram id="mogvwWM-gPOIvBLYb1fR" name="Page-2">ddG9DoIwEADgp+le2xjdEXVxYnBu6EmbFI6UM6BPL6QFbNCp1+/u+stkVg8Xr1pzQw2OCa4HJk9MiN2Oi3GY5BXkuD8EqLzVsWiFwr4hIo/6tBq6pJAQHdk2xRKbBkpKTHmPfVr2QJfu2qoKNlCUym31bjWZ+RZ89SvYytBy4Zip1VwcoTNKY/9FMmcy84gUonrIwE2PN79L6Dv/yS4H89DQj4YxWNceJ8kPyfwD</diagram></mxfile>

fastapi_2fa/api/deps/db.py

@@ -15,3 +15,4 @@ async def get_db() -> Generator:
     finally:
         print('closing db...')
         await db.close()
+        print('..db closed!')

fastapi_2fa/api/deps/users.py

@@ -6,6 +6,7 @@
 
 from fastapi_2fa.api.deps.db import get_db
 from fastapi_2fa.core.config import settings
+from fastapi_2fa.core.two_factor_auth import verify_token
 from fastapi_2fa.crud.users import user_crud
 from fastapi_2fa.models.users import User
 from fastapi_2fa.schemas.jwt_token_schema import JwtTokenPayload
@@ -69,3 +70,17 @@ async def get_authenticated_user_pre_tfa(
         "and validate TFA token within "
         f"{settings.PRE_TFA_TOKEN_EXPIRE_MINUTES} minutes",
     )
+
+
+async def get_authenticated_tfa_user(
+    tfa_token: str,
+    db: Session = Depends(get_db),
+    user: User = Depends(get_authenticated_user_pre_tfa)
+) -> User:
+    if verify_token(user=user, token=tfa_token):
+        return user
+
+    raise HTTPException(
+        status_code=status.HTTP_403_FORBIDDEN,
+        detail="TOTP token mismatch"
+    )

fastapi_2fa/api/endpoints/api_v1/auth.py

@@ -42,12 +42,10 @@ async def signup(
     db: Session = Depends(get_db), user_data: UserCreate = Depends()
 ) -> Any:
     try:
-        # async with db.begin_nested():
         user = await user_crud(transaction=True).create(
             db=db,
             user=user_data,
         )
-        # raise Exception
         if user_data.tfa_enabled:
             device, qr_code = await device_crud(transaction=True).create(
                 db=db,
@@ -78,6 +76,17 @@ async def signup(
     summary="Create access and refresh tokens for user",
     status_code=status.HTTP_200_OK,
     response_model=JwtTokenSchema | PreTfaJwtTokenSchema,
+    responses={
+        200: {
+            "description": "Credentials ok and TFA disabled",
+        },
+        202: {
+            "description": "Credentials ok and TFA enabled",
+        },
+        401: {
+            "description": "Invalid credentials",
+        },
+    }
 )
 async def login(
     response: Response,
@@ -91,7 +100,7 @@ async def login(
     # wrong credentials
     if not user:
         raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST,
+            status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Incorrect email or password",
         )
 
@@ -115,7 +124,9 @@ async def login(
 
 
 @auth_router.get(
-    "/test-token", summary="Test if the access token is ok", response_model=UserOut
+    "/test-token",
+    summary="Authenticated endpoint to test if access token is ok",
+    response_model=UserOut
 )
 async def test_token(user: User = Depends(get_authenticated_user)):
     """
@@ -124,7 +135,11 @@ async def test_token(user: User = Depends(get_authenticated_user)):
     return user
 
 
-@auth_router.post("/refresh", summary="Refresh token", response_model=JwtTokenSchema)
+@auth_router.post(
+    "/refresh",
+    summary="Refresh token for elapsed access token",
+    response_model=JwtTokenSchema
+)
 async def refresh_token(
     db: Session = Depends(get_db), refresh_token: str = Body(embed=True, title='refresh token')
 ):

fastapi_2fa/api/endpoints/api_v1/two_factor_auth.py

@@ -5,13 +5,13 @@
 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.api.deps.users import (get_authenticated_tfa_user,
+                                        get_authenticated_user,
                                         get_authenticated_user_pre_tfa)
 from fastapi_2fa.core import security
 from fastapi_2fa.core.enums import DeviceTypeEnum
 from fastapi_2fa.core.two_factor_auth import (qr_code_from_key,
-                                              verify_backup_token,
-                                              verify_token)
+                                              verify_backup_token)
 from fastapi_2fa.core.utils import send_mail_backup_tokens
 from fastapi_2fa.crud.backup_token import backup_token_crud
 from fastapi_2fa.crud.device import device_crud
@@ -32,17 +32,11 @@
 async def login_tfa(
     tfa_token: str,
     db: Session = Depends(get_db),
-    user: User = Depends(get_authenticated_user_pre_tfa),
+    user: User = Depends(get_authenticated_tfa_user),
 ) -> Any:
-    if verify_token(user=user, token=tfa_token):
-        return JwtTokenSchema(
-            access_token=security.create_jwt_access_token(user.id),
-            refresh_token=security.create_jwt_refresh_token(user.id),
-        )
-
-    raise HTTPException(
-        status_code=status.HTTP_403_FORBIDDEN,
-        detail="TOTP token mismatch"
+    return JwtTokenSchema(
+        access_token=security.create_jwt_access_token(user.id),
+        refresh_token=security.create_jwt_refresh_token(user.id),
     )
 
 

fastapi_2fa/core/utils.py

@@ -31,20 +31,26 @@ def __repr__(self) -> str:
 
 
 def send_mail_backup_tokens(user: User, device: Device) -> AsyncResult:
-    email_ob = Email(
+    email_obj = Email(
         to_=[user.email],
         from_=settings.FAKE_EMAIL_SENDER,
         text_=f"Backup tokens : {device.backup_tokens}"
     )
-    task = send_email_task.apply_async(kwargs={'email': email_ob.to_json()})
+    print(
+        f"Sending email task to celery, email:\n\n***\n{email_obj.text_}\n***\n"
+    )
+    task = send_email_task.apply_async(kwargs={'email': email_obj.to_json()})
     return task
 
 
 def send_mail_totp_token(user: User, token: str) -> AsyncResult:
-    email_ob = Email(
+    email_obj = Email(
         to_=[user.email],
         from_=settings.FAKE_EMAIL_SENDER,
         text_=f"Access TOTP token : {token}"
     )
-    task = send_email_task.apply_async(kwargs={'email': email_ob.to_json()})
+    print(
+        f"Sending email task to celery, email:\n\n***\n{email_obj.text_}\n***\n"
+    )
+    task = send_email_task.apply_async(kwargs={'email': email_obj.to_json()})
     return task

fastapi_2fa/crud/users.py

@@ -29,6 +29,7 @@ async def authenticate(db: Session, email: EmailStr, password: str) -> User | No
         user: User = await UserCrud.get_by_email(db=db, email=email)
         if not user:
             return None
+        user = user[0]
         if not verify_password(password=password, hashed_password=user.hashed_password):
             return None
         return user
@@ -50,7 +51,7 @@ async def update(
     async def get_by_email(db: Session, email: EmailStr) -> User | None:
         query = select(User).where(User.email == email)
         model_query = await db.execute(query)
-        (model_instance,) = model_query.first()
+        model_instance = model_query.first()
         return model_instance
 
 

fastapi_2fa/main.py

@@ -4,6 +4,8 @@
 
 from fastapi_2fa.api.endpoints.router import router
 from fastapi_2fa.core.config import settings
+from fastapi_2fa.db.session import SessionLocal
+from fastapi_2fa.tasks.celery_conf import celery
 
 app = FastAPI(
     title=settings.PROJECT_NAME,
@@ -15,6 +17,33 @@
 app.include_router(router=router, prefix=settings.API_V1_STR)
 
 
+@app.on_event("startup")
+async def app_startup():
+    # CHECK DB
+    try:
+        db = SessionLocal()
+        await db.execute("SELECT 1")
+        print("Database is connected")
+    except Exception:
+        raise RuntimeError("Problems reaching db")
+    finally:
+        await db.close()
+
+    # CHECK CELERY BROKER
+    try:
+        celery.broker_connection().ensure_connection(max_retries=3)
+    except Exception as ex:
+        raise RuntimeError("Failed to connect to celery broker, {}".format(str(ex)))
+    print('Celery broker is running')
+
+    # CHECK CELERY WORKER
+    insp = celery.control.inspect()
+    availability = insp.ping()
+    if not availability:
+        raise RuntimeError('Celery worker is not running')
+    print('Celery worker is running')
+
+
 # Set all CORS enabled origins
 if settings.BACKEND_CORS_ORIGINS:
     app.add_middleware(