202 lines
7.3 KiB
Python
202 lines
7.3 KiB
Python
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"}
|