384 lines
11 KiB
Python

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