from datetime import timedelta from typing import Any from fastapi import APIRouter, Depends, HTTPException, status from fastapi.security import OAuth2PasswordRequestForm from sqlalchemy.orm import Session from pydantic import ValidationError from app import crud, models, schemas from app.api import deps from app.core import security from app.core.config import settings from app.core.email import send_email_verification, send_password_reset from app.services.two_factor import generate_totp_secret, generate_qr_code, verify_totp router = APIRouter() @router.post("/register", response_model=schemas.User) def register( *, db: Session = Depends(deps.get_db), user_in: schemas.UserCreate, ) -> Any: """ Register a new user. """ user = crud.user.get_by_email(db, email=user_in.email) if user: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Email already registered", ) user = crud.user.create(db, obj_in=user_in) # Create wallets for the user from app.models.wallet import WalletType crud.wallet.create_for_user(db, user_id=user.id, wallet_type=WalletType.SPOT) crud.wallet.create_for_user(db, user_id=user.id, wallet_type=WalletType.TRADING) # Generate and send verification email token = security.create_email_verification_token(user.email) send_email_verification(user.email, token) return user @router.post("/login", response_model=schemas.Token) def login( db: Session = Depends(deps.get_db), form_data: OAuth2PasswordRequestForm = Depends(), ) -> Any: """ Get an access token for future requests. """ user = crud.user.authenticate( db, email=form_data.username, password=form_data.password ) if not user: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect email or password", headers={"WWW-Authenticate": "Bearer"}, ) if not crud.user.is_active(user): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Inactive user", ) # If 2FA is enabled, return a special token requiring 2FA if user.is_two_factor_enabled: temp_token = security.create_access_token( subject=str(user.id), role=user.role, expires_delta=timedelta(minutes=15) ) return { "access_token": temp_token, "token_type": "bearer", "requires_two_factor": True } # Regular login flow access_token = security.create_access_token( subject=str(user.id), role=user.role, ) refresh_token = security.create_refresh_token( subject=str(user.id), role=user.role, ) return { "access_token": access_token, "token_type": "bearer", "refresh_token": refresh_token, "requires_two_factor": False } @router.post("/login/2fa", response_model=schemas.Token) def login_two_factor( *, db: Session = Depends(deps.get_db), two_factor_data: schemas.TwoFactorLogin, ) -> Any: """ Complete login with 2FA verification. """ try: payload = security.jwt.decode( two_factor_data.token, settings.JWT_SECRET_KEY, algorithms=[settings.ALGORITHM] ) user_id = payload.get("sub") except (security.jwt.JWTError, ValidationError): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token", headers={"WWW-Authenticate": "Bearer"}, ) user = crud.user.get(db, id=user_id) if not user or not user.is_two_factor_enabled: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found or 2FA not enabled", ) if not verify_totp(user.two_factor_secret, two_factor_data.code): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid 2FA code", ) # Generate full access tokens after 2FA validation access_token = security.create_access_token( subject=str(user.id), role=user.role, ) refresh_token = security.create_refresh_token( subject=str(user.id), role=user.role, ) return { "access_token": access_token, "token_type": "bearer", "refresh_token": refresh_token, "requires_two_factor": False } @router.post("/refresh-token", response_model=schemas.Token) def refresh_token( *, db: Session = Depends(deps.get_db), refresh_token_in: schemas.RefreshToken, ) -> Any: """ Refresh access token. """ try: payload = security.jwt.decode( refresh_token_in.refresh_token, settings.JWT_SECRET_KEY, algorithms=[settings.ALGORITHM] ) user_id = payload.get("sub") except (security.jwt.JWTError, ValidationError): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token", headers={"WWW-Authenticate": "Bearer"}, ) user = crud.user.get(db, id=user_id) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="User not found", ) if not crud.user.is_active(user): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Inactive user", ) access_token = security.create_access_token( subject=str(user.id), role=user.role, ) return { "access_token": access_token, "token_type": "bearer", "refresh_token": refresh_token_in.refresh_token, "requires_two_factor": False } @router.post("/verify-email") def verify_email( *, db: Session = Depends(deps.get_db), email_verification: schemas.EmailVerification, ) -> Any: """ Verify user email. """ email = security.verify_token(email_verification.token, "email_verification") if not email: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid or expired token", ) user = crud.user.get_by_email(db, email=email) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="User not found", ) if crud.user.is_verified(user): return {"message": "Email already verified"} crud.user.set_verified(db, user=user) return {"message": "Email successfully verified"} @router.post("/request-password-reset") def request_password_reset( *, db: Session = Depends(deps.get_db), password_reset: schemas.PasswordReset, ) -> Any: """ Request password reset. """ user = crud.user.get_by_email(db, email=password_reset.email) if not user: # Don't reveal that the user doesn't exist return {"message": "If your email is registered, you will receive a password reset link"} token = security.create_password_reset_token(user.email) send_password_reset(user.email, token) return {"message": "Password reset email sent"} @router.post("/reset-password") def reset_password( *, db: Session = Depends(deps.get_db), password_reset: schemas.PasswordResetConfirm, ) -> Any: """ Reset password. """ email = security.verify_token(password_reset.token, "password_reset") if not email: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid or expired token", ) user = crud.user.get_by_email(db, email=email) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="User not found", ) crud.user.update(db, db_obj=user, obj_in={"password": password_reset.new_password}) return {"message": "Password successfully reset"} @router.post("/enable-2fa", response_model=dict) def enable_two_factor( *, db: Session = Depends(deps.get_db), current_user: models.User = Depends(deps.get_current_active_verified_user), two_factor_setup: schemas.TwoFactorSetup, ) -> Any: """ Enable 2FA for the user. """ # Verify user's password before enabling 2FA if not security.verify_password(two_factor_setup.password, current_user.hashed_password): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect password", ) # Generate 2FA secret secret = generate_totp_secret() # Generate QR code qr_code = generate_qr_code(secret, current_user.email) # Save the 2FA secret (not enabled yet until user verifies) crud.user.set_user_two_factor_secret(db, user=current_user, secret=secret) return { "secret": secret, "qr_code": qr_code, "message": "2FA setup initialized. Verify with the code to enable 2FA." } @router.post("/verify-2fa") def verify_two_factor( *, db: Session = Depends(deps.get_db), current_user: models.User = Depends(deps.get_current_active_verified_user), two_factor_verify: schemas.TwoFactorVerify, ) -> Any: """ Verify and enable 2FA. """ if not current_user.two_factor_secret: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="2FA not initialized", ) if verify_totp(current_user.two_factor_secret, two_factor_verify.code): crud.user.enable_two_factor(db, user=current_user) return {"message": "2FA successfully enabled"} raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid verification code", ) @router.post("/disable-2fa") def disable_two_factor( *, db: Session = Depends(deps.get_db), current_user: models.User = Depends(deps.get_current_active_verified_user), two_factor_disable: schemas.TwoFactorDisable, ) -> Any: """ Disable 2FA. """ if not current_user.is_two_factor_enabled: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="2FA not enabled", ) # Verify user's password before disabling 2FA if not security.verify_password(two_factor_disable.password, current_user.hashed_password): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect password", ) # Verify 2FA code before disabling if not verify_totp(current_user.two_factor_secret, two_factor_disable.code): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid 2FA code", ) crud.user.disable_two_factor(db, user=current_user) return {"message": "2FA successfully disabled"} @router.get("/me", response_model=schemas.User) def read_user_me( current_user: models.User = Depends(deps.get_current_active_verified_user), ) -> Any: """ Get current user. """ return current_user