claudia-docs-api/app/routers/auth.py

from datetime import datetime

from fastapi import APIRouter, Depends, HTTPException, Request, Response
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import get_db
from app.models.agent import Agent
from app.schemas.auth import AgentCreate, AgentLogin, AgentResponse, RefreshResponse, TokenResponse
from app.services import auth as auth_service

router = APIRouter(prefix="/api/v1/auth", tags=["auth"])

COOKIE_NAME = "refresh_token"
COOKIE_PATH = "/api/v1/auth/refresh"
COOKIE_MAX_AGE = 60 * 60 * 24 * 7  # 7 days


def _set_refresh_cookie(response: Response, token: str):
    response.set_cookie(
        key=COOKIE_NAME,
        value=token,
        max_age=COOKIE_MAX_AGE,
        httponly=True,
        secure=True,  # True in production with HTTPS
        samesite="lax",
        path=COOKIE_PATH,
    )


def _clear_refresh_cookie(response: Response):
    response.set_cookie(
        key=COOKIE_NAME,
        value="",
        max_age=0,
        httponly=True,
        secure=True,  # True in production with HTTPS
        samesite="lax",
        path=COOKIE_PATH,
    )


async def get_current_agent(request: Request, db: AsyncSession) -> Agent:
    """Get the current authenticated agent from request."""
    auth_header = request.headers.get("authorization", "")
    if not auth_header.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Not authenticated")

    token = auth_header[7:]
    payload = auth_service.decode_token(token)
    if not payload:
        raise HTTPException(status_code=401, detail="Invalid token")

    jti = payload.get("jti")
    if jti:
        is_blocked = await auth_service.is_token_blocklisted(db, jti)
        if is_blocked:
            raise HTTPException(status_code=401, detail="Token revoked")

    agent_id = payload.get("sub")
    if not agent_id:
        raise HTTPException(status_code=401, detail="Invalid token payload")

    agent = await auth_service.get_agent_by_id(db, agent_id)
    if not agent:
        raise HTTPException(status_code=401, detail="Agent not found")

    return agent


@router.post("/register", response_model=AgentResponse, status_code=201)
async def register(payload: AgentCreate, db: AsyncSession = Depends(get_db)):
    existing = await auth_service.get_agent_by_username(db, payload.username)
    if existing:
        raise HTTPException(status_code=400, detail="Username already exists")

    agent = await auth_service.create_agent(db, payload.username, payload.password)
    return AgentResponse.model_validate(agent)


@router.post("/login", response_model=TokenResponse)
async def login(
    payload: AgentLogin,
    request: Request,
    response: Response,
    db: AsyncSession = Depends(get_db),
):
    agent = await auth_service.get_agent_by_username(db, payload.username)
    if not agent or not auth_service.verify_password(payload.password, agent.password_hash):
        raise HTTPException(status_code=401, detail="Invalid credentials")

    access_token, jti = auth_service.create_access_token(agent.id, agent.role)

    refresh_token = auth_service.create_refresh_token()
    user_agent = request.headers.get("user-agent")
    client_ip = request.client.host if request.client else None
    await auth_service.save_refresh_token(db, agent.id, refresh_token, user_agent, client_ip)

    _set_refresh_cookie(response, refresh_token)
    return TokenResponse(access_token=access_token)


@router.get("/me", response_model=AgentResponse)
async def get_me(request: Request, db: AsyncSession = Depends(get_db)):
    agent = await get_current_agent(request, db)
    return AgentResponse.model_validate(agent)


@router.post("/refresh", response_model=RefreshResponse)
async def refresh(request: Request, response: Response, db: AsyncSession = Depends(get_db)):
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        raise HTTPException(status_code=401, detail="No refresh token")

    user_agent = request.headers.get("user-agent")
    client_ip = request.client.host if request.client else None

    result = await auth_service.rotate_refresh_token(db, token, user_agent, client_ip)
    if not result:
        raise HTTPException(status_code=401, detail="Invalid or expired refresh token")

    new_rt, new_token = result

    agent = await auth_service.get_agent_by_id(db, new_rt.user_id)
    if not agent:
        raise HTTPException(status_code=401, detail="Agent not found")

    access_token, jti = auth_service.create_access_token(agent.id, agent.role)
    _set_refresh_cookie(response, new_token)
    return RefreshResponse(access_token=access_token)


@router.post("/logout", status_code=204)
async def logout(request: Request, response: Response, db: AsyncSession = Depends(get_db)):
    token = request.cookies.get(COOKIE_NAME)
    if token:
        await auth_service.revoke_refresh_token(db, token)

    auth_header = request.headers.get("authorization", "")
    if auth_header.startswith("Bearer "):
        access_token = auth_header[7:]
        payload = auth_service.decode_token(access_token)
        if payload and "jti" in payload:
            exp = datetime.utcfromtimestamp(payload["exp"])
            await auth_service.add_to_blocklist(db, payload["jti"], exp)

    _clear_refresh_cookie(response)
    return Response(status_code=204)


@router.post("/logout-all", status_code=204)
async def logout_all(request: Request, response: Response, db: AsyncSession = Depends(get_db)):
    auth_header = request.headers.get("authorization", "")
    agent = None
    if auth_header.startswith("Bearer "):
        access_token = auth_header[7:]
        payload = auth_service.decode_token(access_token)
        if payload and "sub" in payload:
            agent = await auth_service.get_agent_by_id(db, payload["sub"])

    if agent:
        await auth_service.revoke_all_user_tokens(db, agent.id)
        if payload and "jti" in payload:
            exp = datetime.utcfromtimestamp(payload["exp"])
            await auth_service.add_to_blocklist(db, payload["jti"], exp)

    _clear_refresh_cookie(response)
    return Response(status_code=204)