384 lines
11 KiB
Python
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 |