from fastapi import FastAPI, Depends, HTTPException, status, APIRouter from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm from sqlalchemy.orm import Session from app.core import security from app.api import schemas from app.db import models from jose import JWTError, jwt from app.core.security import get_current_user import pyotp import qrcode import base64 from io import BytesIO from fastapi.responses import StreamingResponse # бд def get_db(): db = models.SessionLocal() try: yield db finally: db.close() authRouter = APIRouter( prefix="/auth", tags=[], ) # регистрация @authRouter.post("/register") async def register(username: str, password: str, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)): if current_user.id != 1: raise HTTPException( status_code=403, detail='Forbidden' ) if len(password.encode('utf-8')) > 72: raise HTTPException( status_code=400, detail="Пароль слишком длинный (макс. 72 байта)") db_user = db.query(models.User).filter( models.User.username == username).first() if db_user: raise HTTPException( status_code=400, detail="Пользователь уже существует") hashed_pwd = security.get_password_hash(password) new_user = models.User(username=username, hashed_password=hashed_pwd) db.add(new_user) db.commit() return {"status": "ok", "message": "User created", "id": new_user.id} @authRouter.post("/hash") async def register(password: str): if len(password.encode('utf-8')) > 72: raise HTTPException( status_code=400, detail="Пароль слишком длинный (макс. 72 байта)") hashed_pwd = security.get_password_hash(password) return {"password": hashed_pwd} # вход @authRouter.post("/login") async def login(data: schemas.LoginRequest, db: Session = Depends(get_db)): print(f"Login attempt: username={data.username}, totp_code provided={bool(data.totp_code)}") user = db.query(models.User).filter( models.User.username == data.username).first() if not user or not security.verify_password(data.password, user.hashed_password): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Неверный логин или пароль", headers={"WWW-Authenticate": "Bearer"}, ) if user.totp_secret: if not data.totp_code: raise HTTPException(status_code=400, detail="TOTP код требуется") totp = pyotp.TOTP(user.totp_secret) if not totp.verify(data.totp_code): raise HTTPException(status_code=400, detail="Неверный TOTP код") access_token = security.create_access_token(data={"sub": str(user.id)}) refresh_token = security.create_refresh_token(data={"sub": str(user.id)}) return { "access_token": access_token, "refresh_token": refresh_token, "token_type": "bearer", "user_id": user.id } @authRouter.post("/totp/enable") async def enable_totp(current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)): # Загружаем свежую копию user из БД user = db.query(models.User).filter(models.User.id == current_user.id).first() if not user: raise HTTPException(status_code=400, detail="Пользователь не найден") #if user.totp_secret: #raise HTTPException(status_code=400, detail="TOTP уже включен") secret = pyotp.random_base32() user.totp_temp_secret = secret db.commit() print(f"TOTP enabled for user {user.id}, secret saved") # Генерировать QR totp = pyotp.TOTP(secret) uri = totp.provisioning_uri(name=user.username, issuer_name="Chepuhagram") img = qrcode.make(uri) buf = BytesIO() img.save(buf, format='PNG') buf.seek(0) qr_base64 = base64.b64encode(buf.getvalue()).decode('utf-8') qr_data_url = f"data:image/png;base64,{qr_base64}" return {"secret": secret, "qr_code": qr_data_url} @authRouter.post("/totp/verify") async def verify_totp(data: schemas.TOTPVerifyRequest, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)): # Загружаем свежую копию user из БД user = db.query(models.User).filter(models.User.id == current_user.id).first() if not user: raise HTTPException(status_code=400, detail="Пользователь не найден") if not user.totp_temp_secret: raise HTTPException(status_code=400, detail="TOTP не включен") try: totp = pyotp.TOTP(user.totp_temp_secret) code_str = str(data.code).strip() is_valid = totp.verify(code_str) print(f"TOTP verify: user_id={user.id}, code={code_str}, secret_set={bool(user.totp_temp_secret)}, valid={is_valid}") if is_valid: user.totp_secret = user.totp_temp_secret user.totp_temp_secret = None db.commit() return {"status": "ok"} else: raise HTTPException(status_code=400, detail="Неверный код") except HTTPException: raise except Exception as e: print(f"TOTP verify error: {str(e)}") raise HTTPException(status_code=500, detail=f"Ошибка верификации: {str(e)}") @authRouter.post("/totp/disable") async def disable_totp(current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)): user = db.query(models.User).filter(models.User.id == current_user.id).first() if user: user.totp_secret = None db.commit() return {"status": "ok"} @authRouter.post("/refresh") async def refresh_token(data: schemas.RefreshRequest): try: payload = jwt.decode(data.refresh_token, security.SECRET_KEY, algorithms=[ security.ALGORITHM]) user_id = str(payload.get("sub")) if user_id is None: raise HTTPException(status_code=401) new_access_token = security.create_access_token(data={"sub": user_id}) new_refresh_token = security.create_refresh_token( data={"sub": user_id}) return {"refresh_token": new_refresh_token, "access_token": new_access_token, "token_type": "bearer"} except JWTError: raise HTTPException(status_code=401, detail="Refresh token expired") @authRouter.post("/setup-account") async def setup_account(data: schemas.SetupAccount, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)): user_to_update = db.merge(current_user) user_to_update.first_name = data.first_name user_to_update.last_name = data.last_name user_to_update.public_key = data.public_key user_to_update.encrypted_private_key = data.encrypted_private_key db.commit() db.refresh(user_to_update) return {"status": "ok", "message": "Account setup completed"} @authRouter.post("/update-fcm") async def update_fcm(token: str, current_user: models.User = Depends(get_current_user), db: Session = Depends(get_db)): user_to_update = db.merge(current_user) user_to_update.fcm_token = token db.commit() db.refresh(user_to_update) return {"status": "ok"}