FabledGrimoire/src/grimoire/services/oauth.py

import logging
import urllib.parse

import httpx

from grimoire.config import Config
from grimoire.models import async_session
from grimoire.models.user import User
from sqlalchemy import select

logger = logging.getLogger(__name__)

_oidc_config: dict | None = None


async def get_oidc_config() -> dict:
    global _oidc_config
    if _oidc_config is not None:
        return _oidc_config

    discovery_url = Config.OIDC_ISSUER.rstrip("/") + "/.well-known/openid-configuration"
    async with httpx.AsyncClient() as client:
        resp = await client.get(discovery_url, follow_redirects=True)
        resp.raise_for_status()
        _oidc_config = resp.json()
    logger.info("OIDC config loaded from %s", discovery_url)
    return _oidc_config


async def build_auth_url(state: str, code_challenge: str) -> str:
    oidc = await get_oidc_config()
    redirect_uri = Config.BASE_URL.rstrip("/") + "/api/auth/oauth/callback"
    params = {
        "response_type": "code",
        "client_id": Config.OIDC_CLIENT_ID,
        "redirect_uri": redirect_uri,
        "scope": Config.OIDC_SCOPES,
        "state": state,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
    }
    return oidc["authorization_endpoint"] + "?" + urllib.parse.urlencode(params)


async def exchange_code(code: str, redirect_uri: str, code_verifier: str) -> dict:
    oidc = await get_oidc_config()
    async with httpx.AsyncClient() as client:
        resp = await client.post(
            oidc["token_endpoint"],
            data={
                "grant_type": "authorization_code",
                "code": code,
                "redirect_uri": redirect_uri,
                "client_id": Config.OIDC_CLIENT_ID,
                "client_secret": Config.OIDC_CLIENT_SECRET,
                "code_verifier": code_verifier,
            },
        )
        resp.raise_for_status()
        return resp.json()


async def get_userinfo(access_token: str) -> dict:
    oidc = await get_oidc_config()
    async with httpx.AsyncClient() as client:
        resp = await client.get(
            oidc["userinfo_endpoint"],
            headers={"Authorization": f"Bearer {access_token}"},
        )
        resp.raise_for_status()
        return resp.json()


async def find_or_create_oauth_user(
    sub: str, email: str, preferred_username: str
) -> User:
    async with async_session() as session:
        # 1. Look up by oauth_sub
        result = await session.execute(select(User).where(User.oauth_sub == sub))
        user = result.scalars().first()
        if user:
            return user

        # 2. Look up by email — auto-link
        if email:
            from sqlalchemy import func
            result = await session.execute(
                select(User).where(func.lower(User.email) == email.lower())
            )
            user = result.scalars().first()
            if user:
                user.oauth_sub = sub
                await session.commit()
                await session.refresh(user)
                logger.info(
                    "Linked OAuth sub %s to existing user %d (%s) via email",
                    sub, user.id, user.username,
                )
                return user

        # 3. Create new user — pick a non-colliding username
        base_username = preferred_username or (email.split("@")[0] if email else "user")
        candidate = base_username
        suffix = 2
        while True:
            result = await session.execute(
                select(User).where(User.username == candidate)
            )
            if not result.scalars().first():
                break
            candidate = f"{base_username}_{suffix}"
            suffix += 1

        user = User(
            username=candidate,
            email=email or None,
            password_hash=None,
            oauth_sub=sub,
        )
        session.add(user)
        await session.commit()
        await session.refresh(user)
        logger.info("Auto-provisioned OAuth user %d (%s) for sub %s", user.id, user.username, sub)
        return user