Spaces:
Build error
Build error
File size: 4,408 Bytes
b7a7f32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
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!"}
|