app/api/v1/endpoints/auth.py - master - euorik/devAPP

app/api/v1/endpoints/auth.py
280 строк10 KB
euorik
Добавил эндпоинт Auth
01 авг 2026, 02:50
01 авг 2026, 02:50025274b
import os
import json
import secrets
from datetime import datetime, timedelta
from typing import Optional

import redis.asyncio as redis_async
from fastapi import APIRouter, HTTPException, Depends, Header, Cookie, Query, status
from fastapi.responses import JSONResponse
from jose import JWTError, jwt
from passlib.context import CryptContext
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.exc import IntegrityError

# --- Конфигурация из переменных окружения ---
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+asyncpg://user:password@localhost:5432/auth_db")
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
SECRET_KEY = os.getenv("SECRET_KEY", secrets.token_urlsafe(32))
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "30"))

# --- SQLAlchemy Base ---
class Base(DeclarativeBase):
    pass

class User(Base):
    __tablename__ = "users"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    username: Mapped[str] = mapped_column(unique=True, index=True)
    email: Mapped[str] = mapped_column(unique=True, index=True)
    password: Mapped[str]
    created_at: Mapped[datetime] = mapped_column(default=datetime.utcnow)

# --- Глобальные переменные (лениво инициализируются) ---
_engine = None
_async_session_factory = None
_redis_client: Optional[redis_async.Redis] = None
_tables_created = False

# --- Функции инициализации (вызываются один раз при первом запросе) ---
async def ensure_engine():
    global _engine, _async_session_factory, _tables_created
    if _engine is None:
        _engine = create_async_engine(DATABASE_URL, echo=True)
        _async_session_factory = async_sessionmaker(_engine, class_=AsyncSession, expire_on_commit=False)
        # Создание таблиц (можно вынести в миграции)
        async with _engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        _tables_created = True
    return _engine

async def get_session() -> AsyncSession:
    await ensure_engine()
    async with _async_session_factory() as session:
        yield session

async def get_redis() -> redis_async.Redis:
    global _redis_client
    if _redis_client is None:
        _redis_client = redis_async.from_url(REDIS_URL, decode_responses=True)
    return _redis_client

# --- Вспомогательные функции ---
pwd_context = CryptContext(schemes=["argon2"], deprecated="auto")

def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password: str) -> str:
    return pwd_context.hash(password)

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None) -> str:
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

def decode_token(token: str) -> dict:
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        email: str = payload.get("sub")
        if email is None:
            raise HTTPException(status_code=401, detail="Invalid token (no subject)")
        return payload
    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid or expired token")

async def is_token_blacklisted(token: str) -> bool:
    r = await get_redis()
    return await r.exists(f"blacklist:{token}")

async def blacklist_token(token: str, ttl: int = 3600):
    r = await get_redis()
    await r.setex(f"blacklist:{token}", ttl, "1")

async def get_cached_user(email: str) -> Optional[dict]:
    r = await get_redis()
    data = await r.get(f"user:{email}")
    return json.loads(data) if data else None

async def cache_user(email: str, user_data: dict, ttl: int = 300):
    r = await get_redis()
    await r.setex(f"user:{email}", ttl, json.dumps(user_data))

# --- Pydantic схемы ---
class UserRegister(BaseModel):
    username: str
    email: str
    password: str

class UserLogin(BaseModel):
    email: str
    password: str

class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"

class UserOut(BaseModel):
    id: int
    username: str
    email: str
    created_at: datetime

# --- Роутер ---
auth_rout = APIRouter(tags=["Auth"])

# ============ Endpoints ============

@auth_rout.post("/register", response_model=UserOut, status_code=status.HTTP_201_CREATED)
async def register(user: UserRegister, session: AsyncSession = Depends(get_session)):
    existing = await session.execute(
        select(User).where((User.email == user.email) | (User.username == user.username))
    )
    if existing.scalar_one_or_none():
        raise HTTPException(status_code=400, detail="Email or username already registered")
    
    hashed = get_password_hash(user.password)
    new_user = User(
        username=user.username,
        email=user.email,
        password=hashed,
        created_at=datetime.utcnow()
    )
    session.add(new_user)
    await session.commit()
    await session.refresh(new_user)
    
    await cache_user(new_user.email, {
        "id": new_user.id,
        "username": new_user.username,
        "email": new_user.email,
        "created_at": new_user.created_at.isoformat()
    })
    
    return UserOut(
        id=new_user.id,
        username=new_user.username,
        email=new_user.email,
        created_at=new_user.created_at
    )

@auth_rout.post("/login/header", response_model=TokenResponse)
async def login_header(user: UserLogin, session: AsyncSession = Depends(get_session)):
    result = await session.execute(select(User).where(User.email == user.email))
    db_user = result.scalar_one_or_none()
    if not db_user or not verify_password(user.password, db_user.password):
        raise HTTPException(status_code=401, detail="Invalid email or password")
    
    token = create_access_token(data={"sub": user.email, "username": db_user.username})
    r = await get_redis()
    await r.sadd(f"user_tokens:{user.email}", token)
    await r.expire(f"user_tokens:{user.email}", ACCESS_TOKEN_EXPIRE_MINUTES * 60)
    
    return {"access_token": token, "token_type": "bearer"}

@auth_rout.get("/protected/header", response_model=UserOut)
async def protected_header(
    authorization: str = Header(..., description="Bearer <token>"),
    session: AsyncSession = Depends(get_session)
):
    if not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Invalid authorization header format")
    token = authorization.split(" ", 1)[1]
    if await is_token_blacklisted(token):
        raise HTTPException(status_code=401, detail="Token has been revoked")
    payload = decode_token(token)
    email = payload["sub"]
    
    cached_user = await get_cached_user(email)
    if cached_user:
        return UserOut(**cached_user)
    
    result = await session.execute(select(User).where(User.email == email))
    db_user = result.scalar_one_or_none()
    if not db_user:
        raise HTTPException(status_code=404, detail="User not found")
    user_data = {
        "id": db_user.id,
        "username": db_user.username,
        "email": db_user.email,
        "created_at": db_user.created_at.isoformat()
    }
    await cache_user(email, user_data)
    return UserOut(
        id=db_user.id,
        username=db_user.username,
        email=db_user.email,
        created_at=db_user.created_at
    )

@auth_rout.post("/login/cookie")
async def login_cookie(user: UserLogin, session: AsyncSession = Depends(get_session)):
    result = await session.execute(select(User).where(User.email == user.email))
    db_user = result.scalar_one_or_none()
    if not db_user or not verify_password(user.password, db_user.password):
        raise HTTPException(status_code=401, detail="Invalid email or password")
    token = create_access_token(data={"sub": user.email, "username": db_user.username})
    response = JSONResponse(status_code=200, content={"message": "Logged in", "token_set": True})
    response.set_cookie(
        key="access_token",
        value=token,
        httponly=True,
        max_age=ACCESS_TOKEN_EXPIRE_MINUTES * 60,
        expires=ACCESS_TOKEN_EXPIRE_MINUTES * 60,
        samesite="lax",
        secure=False
    )
    return response

@auth_rout.get("/protected/cookie", response_model=UserOut)
async def protected_cookie(
    access_token: Optional[str] = Cookie(None),
    session: AsyncSession = Depends(get_session)
):
    if not access_token:
        raise HTTPException(status_code=401, detail="No token in cookie")
    if await is_token_blacklisted(access_token):
        raise HTTPException(status_code=401, detail="Token has been revoked")
    payload = decode_token(access_token)
    email = payload["sub"]
    cached_user = await get_cached_user(email)
    if cached_user:
        return UserOut(**cached_user)
    result = await session.execute(select(User).where(User.email == email))
    db_user = result.scalar_one_or_none()
    if not db_user:
        raise HTTPException(status_code=404, detail="User not found")
    return UserOut(
        id=db_user.id,
        username=db_user.username,
        email=db_user.email,
        created_at=db_user.created_at
    )

@auth_rout.get("/protected/query", response_model=UserOut)
async def protected_query(
    token: str = Query(..., description="JWT token as query param"),
    session: AsyncSession = Depends(get_session)
):
    if await is_token_blacklisted(token):
        raise HTTPException(status_code=401, detail="Token has been revoked")
    payload = decode_token(token)
    email = payload["sub"]
    cached_user = await get_cached_user(email)
    if cached_user:
        return UserOut(**cached_user)
    result = await session.execute(select(User).where(User.email == email))
    db_user = result.scalar_one_or_none()
    if not db_user:
        raise HTTPException(status_code=404, detail="User not found")
    return UserOut(
        id=db_user.id,
        username=db_user.username,
        email=db_user.email,
        created_at=db_user.created_at
    )