Spaces:
Sleeping
Sleeping
| from fastapi import APIRouter, Depends, HTTPException, Response, Cookie, status | |
| from sqlalchemy import select | |
| from sqlalchemy.ext.asyncio import AsyncSession | |
| from app.database import get_db | |
| from app.models.user import User | |
| from app.schemas.auth import UserCreate, UserLogin, UserResponse, TokenResponse | |
| from app.services.auth import ( | |
| hash_password, verify_password, | |
| create_access_token, create_refresh_token, decode_token, | |
| get_current_user, | |
| ) | |
| from app.config import get_settings | |
| router = APIRouter(prefix="/api/auth", tags=["auth"]) | |
| settings = get_settings() | |
| async def register(data: UserCreate, db: AsyncSession = Depends(get_db)): | |
| # Check existing | |
| existing = await db.execute( | |
| select(User).where((User.email == data.email) | (User.username == data.username)) | |
| ) | |
| if existing.scalar_one_or_none(): | |
| raise HTTPException(status_code=409, detail="Email or username already registered") | |
| user = User( | |
| email=data.email, | |
| username=data.username, | |
| password_hash=hash_password(data.password), | |
| ) | |
| db.add(user) | |
| await db.commit() | |
| await db.refresh(user) | |
| return user | |
| async def login(data: UserLogin, response: Response, db: AsyncSession = Depends(get_db)): | |
| result = await db.execute(select(User).where(User.email == data.email)) | |
| user = result.scalar_one_or_none() | |
| if not user or not verify_password(data.password, user.password_hash): | |
| raise HTTPException(status_code=401, detail="Invalid credentials") | |
| access_token = create_access_token(user.id) | |
| refresh_token = create_refresh_token(user.id) | |
| response.set_cookie( | |
| key="refresh_token", | |
| value=refresh_token, | |
| httponly=True, | |
| secure=False, # Set True in production with HTTPS | |
| samesite="lax", | |
| max_age=settings.refresh_token_expire_days * 86400, | |
| path="/api/auth", | |
| ) | |
| return TokenResponse(access_token=access_token) | |
| async def refresh( | |
| response: Response, | |
| refresh_token: str | None = Cookie(default=None), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| if not refresh_token: | |
| raise HTTPException(status_code=401, detail="No refresh token") | |
| payload = decode_token(refresh_token) | |
| if payload.get("type") != "refresh": | |
| raise HTTPException(status_code=401, detail="Invalid token type") | |
| user_id = int(payload["sub"]) | |
| result = await db.execute(select(User).where(User.id == user_id)) | |
| user = result.scalar_one_or_none() | |
| if not user: | |
| raise HTTPException(status_code=401, detail="User not found") | |
| new_access = create_access_token(user.id) | |
| new_refresh = create_refresh_token(user.id) | |
| response.set_cookie( | |
| key="refresh_token", | |
| value=new_refresh, | |
| httponly=True, | |
| secure=False, | |
| samesite="lax", | |
| max_age=settings.refresh_token_expire_days * 86400, | |
| path="/api/auth", | |
| ) | |
| return TokenResponse(access_token=new_access) | |
| async def logout(response: Response): | |
| response.delete_cookie("refresh_token", path="/api/auth") | |
| return {"message": "Logged out"} | |
| async def me(user: User = Depends(get_current_user)): | |
| return UserResponse.from_user(user) | |
| async def set_pollinations_key( | |
| data: dict, | |
| user: User = Depends(get_current_user), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| """Save the user's BYOP Pollinations API key.""" | |
| key = data.get("api_key", "").strip() | |
| if not key: | |
| raise HTTPException(status_code=400, detail="api_key is required") | |
| user.pollinations_api_key = key | |
| await db.commit() | |
| return {"message": "Pollinations key saved", "has_pollinations_key": True} | |
| async def delete_pollinations_key( | |
| user: User = Depends(get_current_user), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| """Disconnect the user's Pollinations account.""" | |
| user.pollinations_api_key = None | |
| await db.commit() | |
| return {"message": "Pollinations key removed", "has_pollinations_key": False} | |