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, Form 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 from starlette.status import HTTP_401_UNAUTHORIZED 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_passwordless_create_token, authorize_passwordless_token, verify_passwordless_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, ) router = APIRouter() @router.post( "/web/", response_model=Optional[schemas.user.UserLoginReturn], response_model_exclude_none=True, ) async def login_web_session( db: Session = Depends(deps.get_db), *, form_data: schemas.LoginData, request: Request, response: Response, ) -> Any: if not form_data.username: form_data.username = form_data.email user = cruds.crud_user.authenticate( db, email=form_data.username, password=form_data.password ) if not user: raise HTTPException( status_code=401, detail="Error ID: 111" ) # Incorrect email or password elif not user.is_active: raise HTTPException( status_code=401, detail="Error ID: 112") # Inactive user if user.two_fa_secret: temp_token = await create_2fa_temp_token(user, form_data.remember_me) response.set_cookie("temp_session", temp_token, httponly=True) return { "msg": "2FA required before proceeding!", "two_fa_required": True, "user": None, } else: session_token = await create_sesssion_token( user, form_data.remember_me, request ) response.set_cookie("session", session_token, httponly=True) return { "msg": "Logged in successfully!", "user": user, "two_fa_required": False, } @router.get("/password-less/create") async def generate_passwordless_login_token( db: Session = Depends(deps.get_db), ): token = await create_passwordless_create_token() return {"token": token} @router.post("/password-less/authorize") async def authorize_passwordless_login( db: Session = Depends(deps.get_db), current_user: models.User = Depends(deps.get_current_active_user), token: str = Form(None), ): _ = await authorize_passwordless_token(current_user, token) return {"msg": "Success"} @router.post( "/password-less/verify", response_model=Optional[schemas.user.UserLoginReturn], response_model_exclude_none=True, ) async def verify_passwordless_login( response: Response, request: Request, db: Session = Depends(deps.get_db), token: str = Form(None), ): user_id = await verify_passwordless_token(token) user = cruds.crud_user.get_by_id(db, id=user_id) if not user: raise HTTPException( status_code=HTTP_401_UNAUTHORIZED, detail="Invalid user!") session_token = await create_sesssion_token(user, True, request) response.set_cookie("session", session_token, httponly=True) return {"msg": "Logged in successfully!", "user": user, "two_fa_required": False} @router.post("/signup/", response_model=schemas.Msg) async def sign_up( *, db: Session = Depends(deps.get_db), user_in: schemas.UserSignUp, ) -> Any: if not settings.USERS_OPEN_REGISTRATION: raise HTTPException( status_code=403, detail="Error ID: 129", ) # Open user registration is forbidden on this server user = cruds.crud_user.get_by_email(db, email=user_in.email) if user: raise HTTPException( status_code=400, detail="Email is associated with another user!", ) # The user with this username already exists in the system email_host = user_in.email[user_in.email.index("@") + 1:] if email_host not in settings.ALLOWED_EMAIL_HOST: raise HTTPException( status_code=403, # TODO: Reflected XSS test detail=f"Email of host {email_host} not allowed!", ) user = cruds.crud_user.create( db, obj_in=schemas.UserCreate(**user_in.dict(), profile_pic="") ) await send_verification_email(user=user) return schemas.Msg(msg="Success") @router.post("/resend-verification-email/") async def resend_verification_email( email: str, current_user: models.User = Depends(deps.get_current_admin_or_above), db: Session = Depends(deps.get_db), ): user = cruds.crud_user.get_by_email(db=db, email=email) if not user: raise HTTPException(status_code="404", detail="User doesn't exist") await send_verification_email(user) return schemas.Msg(msg="Success") @router.post("/change-password/") async def change_password( current_password: str = Body(...), new_password: str = Body(...), current_user: models.User = Depends(deps.get_current_user), db: Session = Depends(deps.get_db), ) -> Any: user = cruds.crud_user.authenticate( db, email=current_user.email, password=current_password ) if not user: raise HTTPException( status_code=403, detail="Error ID: 111" ) # Incorrect email or password data = schemas.user.PasswordUpdate( password=new_password, ) cruds.crud_user.update(db=db, db_obj=current_user, obj_in=data) @router.get("/active-sessions/", response_model=List[schemas.auth.ActiveSession]) async def get_active_sessions( current_user: models.User = Depends(deps.get_current_user), ) -> Any: active_sessions = json.loads( await redis_session_client.client.get( f"user_{current_user.id}_sessions", encoding="utf-8" ) ) return active_sessions.get("sessions") @router.get("/logout-all-sessions/") async def logout_all_sessions( current_user: models.User = Depends(deps.get_current_user), ) -> Any: active_sessions = json.loads( await redis_session_client.client.get( f"user_{current_user.id}_sessions", encoding="utf-8" ) ) for session in active_sessions.get("sessions"): await redis_session_client.client.expire(session.get("token"), 0) await redis_session_client.client.expire(f"user_{current_user.id}_sessions", 0) resp = JSONResponse({"status": "success"}) resp.delete_cookie("session") return resp @router.post("/password-recovery/", response_model=schemas.Msg) # @throttle.ip_throttle(rate=3, per=1 * 60 * 60) # @throttle.ip_throttle(rate=1, per=20) async def recover_password( request: Request, email: str, db: Session = Depends(deps.get_db), ) -> Any: """ Password Recovery """ user = cruds.crud_user.get_by_email(db, email=email) if not user: raise HTTPException( status_code=404, detail="Error ID: 113", ) # The user with this username does not exist in the system. await send_reset_password_email(user=user) return {"msg": "Password recovery email sent"} @router.post("/reset-password/", response_model=schemas.Msg) async def reset_password( request: Request, token: str = Body(...), new_password: str = Body(...), db: Session = Depends(deps.get_db), ) -> Any: """ Reset password """ uid = await verify_password_reset_token(token) user = cruds.crud_user.get_by_id(db, id=uid) if not user: raise HTTPException( status_code=404, detail="Error ID: 114", ) # The user with this username does not exist in the system. elif not cruds.crud_user.is_active(user): raise HTTPException( status_code=400, detail="Error ID: 115") # Inactive user hashed_password = get_password_hash(new_password) user.hashed_password = hashed_password db.add(user) db.commit() return {"msg": "Password updated successfully"} @router.post("/verify/", response_model=schemas.Msg) async def verify_account( token: str, db: Session = Depends(deps.get_db), ) -> Any: uid = await verify_user_verify_token(token) user = cruds.crud_user.get_by_id(db, id=uid) if not user: raise HTTPException( status_code=404, detail="Error ID: 146", ) # The user with this username does not exist in the system. cruds.crud_user.verify_user(db=db, db_obj=user) return {"msg": "Verified successfully"} @router.get("/logout/", response_model=schemas.Token) async def session_logout( session: str = ReqCookie(None), ) -> Any: if not session: raise HTTPException(status_code=401, detail="Invalid session token!") await expire_web_session(session) resp = JSONResponse({"status": "success"}) resp.delete_cookie("session") return resp @router.get("/thtest1") @throttle.ip_throttle(rate=10, per=60) async def throttle_test( request: Request, db: Session = Depends(deps.get_db), current_user: models.User = Depends(deps.get_current_user), ): return "Throttle test endpoint 1 Hello" @router.get("/thtest2") @throttle.user_throttle(rate=20, per=60) async def throttle_test1( request: Request, db: Session = Depends(deps.get_db), current_user: models.User = Depends(deps.get_current_user), ): return "Throttle test endpoint 2"