api/endpoints/two_fa.py

import json
import os
from typing import Any, List, Optional

import aiofiles
from fastapi import APIRouter, Body
from fastapi import Cookie as ReqCookie
from fastapi import Depends, File, HTTPException, Request, UploadFile
from fastapi.params import Cookie
from sqlalchemy.orm import Session
from sqlalchemy.sql.expression import update
from sqlalchemy.sql.functions import current_user
from starlette.responses import JSONResponse, Response

import cruds
import models
import schemas
from core import throttle
from core.config import settings
from core.db import redis_session_client
from core.security import (
    create_sesssion_token,
    get_password_hash,
    create_2fa_temp_token,
    create_2fa_enable_temp_token
)
from cruds import group
from schemas.user import UserUpdate, VerifyUser
from utils import deps
from utils.utils import (
    expire_web_session,
    generate_password_reset_token,
    send_reset_password_email,
    send_verification_email,
    verify_password_reset_token,
    verify_user_verify_token,
)
import pyotp
from cruds import crud_user


router = APIRouter()


@router.get("/enable/request")
async def two_fa_enable_request(
    db: Session = Depends(deps.get_db),
    *,
    current_user: models.User = Depends(deps.get_current_user),
    request: Request,
    response: Response,
) -> Any:
    if current_user.two_fa_secret != None:
        raise HTTPException(
            status_code=409, detail="2FA is already enabled!"
        )

    totp_secret = pyotp.random_base32()
    await create_2fa_enable_temp_token(current_user, totp_secret)

    totp_url = pyotp.totp.TOTP(totp_secret).provisioning_uri(
        name=current_user.email,
        issuer_name=settings.PROJECT_NAME
    )

    return {"msg": "2FA enable requested!", "uri": totp_url, "secret": totp_secret}


@router.post("/enable/confirm")
async def two_fa_enable_confirm(
    db: Session = Depends(deps.get_db),
    *,
    form_data: schemas.Two_FA_Confirm,
    current_user: models.User = Depends(deps.get_current_user),
) -> Any:
    totp_secret = await redis_session_client.client.get(
        f"two_fa_enable_temp_{current_user.id}"
    )
    totp_secret = totp_secret.decode("utf-8")
    if not totp_secret:
        raise HTTPException(
            status_code=403, detail="Invalid or expired TOTP"
        )

    totp = pyotp.TOTP(totp_secret)
    totp_valid = totp.verify(str(form_data.totp), valid_window=1)

    if totp_valid:
        crud_user.enable_2fa(db, secret=totp_secret, db_obj=current_user)
        await redis_session_client.client.delete(
            f"two_fa_enable_temp_{current_user.id}"
        )
        return {"msg": "2FA successfully enabled!"}
    else:
        return {"msg": "Invalid TOTP!"}


@router.post("/login/confirm", response_model=Optional[schemas.user.UserLoginReturn], response_model_exclude_none=True)
async def two_fa_login_confirm(
    db: Session = Depends(deps.get_db),
    *,
    form_data: schemas.Two_FA_Confirm,
    request: Request,
    response: Response
) -> Any:
    token = request.cookies.get("temp_session")

    if token == None:
        raise HTTPException(
            status_code=403, detail="Invalid token!"
        )

    data = json.loads(await redis_session_client.client.get(
        f"two_fa_temp_{token}",
    ))
    # json.dumps({"user": user.id, "remember_me": remember_me}),

    user = crud_user.get(db, id=data.get("user"))

    totp = pyotp.TOTP(user.two_fa_secret)
    totp_valid = totp.verify(str(form_data.totp), valid_window=1)

    if not totp_valid:
        raise HTTPException(
            status_code=403, detail="Invalid TOTP!"
        )

    session_token = await create_sesssion_token(user, data.get("remember_me"), request)

    response.delete_cookie("temp_session")
    response.set_cookie("session", session_token, httponly=True)

    await redis_session_client.client.delete(f"two_fa_temp_{token}")
    return {"msg": "Logged in successfully!", "user": user, "two_fa_required": None}


@router.delete(
    "/disable",
)
async def two_fa_disable(
    db: Session = Depends(deps.get_db),
    *,
    current_user: models.User = Depends(deps.get_current_user),
) -> Any:
    if current_user.two_fa_secret == None:
        raise HTTPException(
            status_code=409, detail="2FA is already disabled!"
        )
    crud_user.disable_2fa(db, db_obj=current_user)

    return {"msg": "2FA successfully disabled!"}