Merge pull request #212 from prefeitura-rio/development

Development

app/config/__init__.py

@@ -81,7 +81,7 @@ def inject_environment_variables(environment: str):
         f"Injecting {len(secrets)} environment variables from Infisical:")
     for secret in secrets:
         logger.info(
-            f" - {secret.secret_name}: {'*' * len(secret.secret_value)}")
+            f" - {secret.secret_name}: {len(secret.secret_value)} chars")
 
 
 environment = getenv_or_action("ENVIRONMENT", action="warn", default="dev")

app/config/base.py

@@ -6,8 +6,17 @@
 # Logging
 LOG_LEVEL = getenv_or_action("LOG_LEVEL", default="INFO")
 
-# BigQuery Project
+# BigQuery Integration
 BIGQUERY_PROJECT = getenv_or_action("BIGQUERY_PROJECT", action="raise")
+BIGQUERY_PATIENT_HEADER_TABLE_ID = getenv_or_action(
+    "BIGQUERY_PATIENT_HEADER_TABLE_ID", action="raise"
+)
+BIGQUERY_PATIENT_SUMMARY_TABLE_ID = getenv_or_action(
+    "BIGQUERY_PATIENT_SUMMARY_TABLE_ID", action="raise"
+)
+BIGQUERY_PATIENT_ENCOUNTERS_TABLE_ID = getenv_or_action(
+    "BIGQUERY_PATIENT_ENCOUNTERS_TABLE_ID", action="raise"
+)
 
 # JWT configuration
 JWT_SECRET_KEY = getenv_or_action("JWT_SECRET_KEY", default=token_bytes(32).hex())

app/models.py

@@ -277,6 +277,10 @@ class User(Model):
     is_superuser = fields.BooleanField(default=False)
     user_class = fields.CharEnumField(enum_type=UserClassEnum, null=True, default=UserClassEnum.PIPELINE_USER)
     data_source = fields.ForeignKeyField("app.DataSource", related_name="users", null=True)
+    # 2FA
+    secret_key = fields.CharField(max_length=255, null=True)
+    is_2fa_required = fields.BooleanField(default=False)
+    is_2fa_activated = fields.BooleanField(default=False)
     created_at = fields.DatetimeField(auto_now_add=True)
     updated_at = fields.DatetimeField(auto_now=True)
 

app/routers/auth.py

@@ -1,43 +1,127 @@
 # -*- coding: utf-8 -*-
-from datetime import timedelta
+import io
 from typing import Annotated
 
 from fastapi import APIRouter, Depends, HTTPException, status
 from fastapi.security import OAuth2PasswordRequestForm
+from fastapi.responses import StreamingResponse
 
-from app import config
 from app.models import User
 from app.types.pydantic_models import Token
-from app.utils import authenticate_user, create_access_token
+from app.utils import authenticate_user, generate_user_token
+from app.security import TwoFactorAuth
+from app.dependencies import (
+    get_current_frontend_user
+)
 
 
 router = APIRouter(prefix="/auth", tags=["Autenticação"])
 
 
 @router.post("/token")
-async def login_for_access_token(
-    form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
+async def login_without_2fa(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
 ) -> Token:
 
-    user: User = await authenticate_user(form_data.username, form_data.password)
+    user = await authenticate_user(form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    if user.is_2fa_required:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="2FA required. Use the /2fa/login/ endpoint",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    return {
+        "access_token": generate_user_token(user),
+        "token_type": "bearer"
+    }
+
+
+@router.post("/2fa/is-2fa-active/")
+async def is_2fa_active(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+) -> bool:
+    user = await authenticate_user(form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    return user.is_2fa_activated
+
 
+@router.post("/2fa/login/")
+async def login_with_2fa(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+    totp_code: str,
+) -> Token:
+
+    user = await authenticate_user(form_data.username, form_data.password)
     if not user:
         raise HTTPException(
-            status_code = status.HTTP_401_UNAUTHORIZED,
-            detail      = "Incorrect username or password",
-            headers     = {"WWW-Authenticate": "Bearer"},
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
         )
 
-    access_token_expires = timedelta(
-        minutes = config.JWT_ACCESS_TOKEN_EXPIRE_MINUTES
-    )
+    secret_key = await TwoFactorAuth.get_or_create_secret_key(user.id)
+    two_factor_auth = TwoFactorAuth(user.id, secret_key)
 
-    access_token = create_access_token(
-        data            = {"sub": user.username},
-        expires_delta   = access_token_expires
-    )
+    is_valid = two_factor_auth.verify_totp_code(totp_code)
+    if not is_valid:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect OTP",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    if not user.is_2fa_activated:
+        user.is_2fa_activated = True
+        await user.save()
 
     return {
-        "access_token": access_token,
-        "token_type": "bearer"
-    }
+        "access_token": generate_user_token(user),
+        "token_type": "bearer",
+    }
+
+
+@router.post("/2fa/enable/")
+async def enable_2fa(
+    current_user: Annotated[User, Depends(get_current_frontend_user)],
+):
+    secret_key = await TwoFactorAuth.get_or_create_secret_key(current_user.id)
+    two_factor_auth = TwoFactorAuth(current_user.id, secret_key)
+
+    return {
+        "secret_key": two_factor_auth.secret_key
+    }
+
+
+@router.get("/2fa/generate-qrcode/")
+async def generate_qrcode(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+):
+    current_user = await authenticate_user(form_data.username, form_data.password)
+    if not current_user:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    secret_key = await TwoFactorAuth.get_or_create_secret_key(current_user.id)
+    two_factor_auth = TwoFactorAuth(current_user.id, secret_key)
+
+    qr_code = two_factor_auth.qr_code
+    if qr_code is None:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    return StreamingResponse(io.BytesIO(qr_code), media_type="image/png")