app/backend/fluksio/api/routes/oauth.py

"""OAuth 2.1 authorization server, for agents reaching the MCP endpoint.

An agent cannot be handed a password, so it registers itself, sends a human to
the browser to approve it, and exchanges the resulting code for a token. The
parts that carry the security are the ones with the least room for
interpretation:

* the code is single-use, short-lived, and only its hash is stored;
* PKCE is required (S256 only), so a code intercepted on its way back is
  useless without the verifier that started the flow;
* the redirect the browser is finally sent to is the one *registered*, never
  the one asked for;
* refresh tokens rotate, and reusing a spent one revokes the whole line, which
  is how a stolen token gets noticed.

Errors here follow RFC 6749 — ``{"error": ...}`` rather than FastAPI's
``{"detail": ...}`` — because that is what OAuth clients parse.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import re
import secrets
import time
import uuid
from collections import OrderedDict, defaultdict, deque
from datetime import UTC, datetime, timedelta
from typing import Any
from urllib.parse import urlencode, urlparse

from fastapi import APIRouter, Depends, Form, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlmodel import col, select

from fluksio.api.deps import (
    CurrentUser,
    SessionDep,
    get_current_active_superuser,
    get_current_user,
)
from fluksio.core import security
from fluksio.core.config import settings
from fluksio.models import (
    Message,
    OAuthAuthorizationCode,
    OAuthAuthorizeInfo,
    OAuthAuthorizeRequest,
    OAuthAuthorizeResponse,
    OAuthClient,
    OAuthClientInfo,
    OAuthClientRegister,
    OAuthRefreshToken,
    User,
)

router = APIRouter(prefix="/oauth", tags=["oauth"])

#: RFC 7636: 43-128 characters from the unreserved set.
_PKCE_RE = re.compile(r"^[A-Za-z0-9._~-]{43,128}$")
_LOOPBACK = {"localhost", "127.0.0.1", "::1"}
_NO_STORE = {"Cache-Control": "no-store", "Pragma": "no-cache"}


def _error(code: str, description: str, status: int = 400) -> JSONResponse:
    """An OAuth error, in the shape clients expect to read."""
    return JSONResponse(
        status_code=status,
        content={"error": code, "error_description": description},
        headers=_NO_STORE,
    )


# -----------------------------------------------------------------------------
# Rate limiting
#
# Registration is open by necessity, so it is capped per address. In-process is
# enough: the engine is one process, and this is a speed bump, not a boundary.
# -----------------------------------------------------------------------------

#: How many distinct buckets are remembered at once. Entries were pruned
#: *within* a deque and never removed, so one key per source address survived
#: for the life of the process.
_MAX_BUCKETS = 4096
_hits: OrderedDict[str, deque[float]] = OrderedDict()


def _too_many(bucket: str, limit: int, window: float, record: bool = True) -> bool:
    now = time.monotonic()
    seen = _hits.get(bucket)
    if seen is None:
        seen = _hits[bucket] = deque()
        while len(_hits) > _MAX_BUCKETS:
            _hits.popitem(last=False)
    else:
        _hits.move_to_end(bucket)
    while seen and now - seen[0] > window:
        seen.popleft()
    if len(seen) >= limit:
        return True
    if record:
        seen.append(now)
    return False


def _client_ip(request: Request) -> str:
    """The caller's address, or the proxy's if it did not forward one.

    Behind Traefik `request.client.host` is the proxy for every request, which
    made every per-address limit here one global bucket — so one caller could
    lock out everyone. The leftmost forwarded entry is the client, and is what
    the panel pairing screen already reads.
    """
    forwarded = request.headers.get("x-forwarded-for", "").split(",")[0].strip()
    if forwarded:
        return forwarded
    return request.client.host if request.client else "unknown"


def too_many(
    request: Request, action: str, limit: int, window: float, record: bool = True
) -> bool:
    """Per-address rate limit, for anything unauthenticated on any router.

    ``record=False`` asks the question without spending an attempt, for a
    caller that only wants to count the ones that failed.
    """
    return _too_many(f"{action}:{_client_ip(request)}", limit, window, record)


# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------


def _hash(value: str) -> str:
    return hashlib.sha256(value.encode()).hexdigest()


def _valid_redirect_uri(value: str) -> bool:
    """https anywhere, or plain http only on the loopback interface.

    An agent running on someone's laptop listens on 127.0.0.1 and has nowhere
    to get a certificate, which OAuth 2.1 allows for exactly that reason.
    """
    parsed = urlparse(value)
    if parsed.fragment:
        return False
    if parsed.scheme == "https":
        return bool(parsed.hostname)
    if parsed.scheme == "http":
        return parsed.hostname in _LOOPBACK
    return False


def _now() -> datetime:
    return datetime.now(UTC)


def _aware(value: datetime) -> datetime:
    """Postgres hands back naive datetimes; compare them in UTC."""
    return value if value.tzinfo else value.replace(tzinfo=UTC)


def _prune(session: SessionDep) -> None:
    """Drop codes long past the window where a replay could still matter."""
    cutoff = _now() - timedelta(days=1)
    for code in session.exec(
        select(OAuthAuthorizationCode).where(OAuthAuthorizationCode.expires_at < cutoff)
    ).all():
        session.delete(code)


def _issue(
    session: SessionDep, user_id: uuid.UUID, client_id: uuid.UUID, family: uuid.UUID
) -> tuple[dict[str, Any], OAuthRefreshToken]:
    """Mint an access token and the refresh token that will replace it."""
    access = security.create_oauth_access_token(
        user_id,
        client_id,
        timedelta(minutes=settings.MCP_TOKEN_EXPIRE_MINUTES),
    )
    refresh_secret = secrets.token_urlsafe(32)
    refresh = OAuthRefreshToken(
        token_hash=_hash(refresh_secret),
        client_id=client_id,
        user_id=user_id,
        family_id=family,
        expires_at=_now() + timedelta(days=settings.MCP_REFRESH_EXPIRE_DAYS),
    )
    session.add(refresh)
    return {
        "access_token": access,
        "token_type": "Bearer",
        "expires_in": settings.MCP_TOKEN_EXPIRE_MINUTES * 60,
        "refresh_token": refresh_secret,
        "scope": security.MCP_SCOPE,
    }, refresh


# -----------------------------------------------------------------------------
# Registration
# -----------------------------------------------------------------------------


@router.post("/register")
def register_client(
    body: OAuthClientRegister, request: Request, session: SessionDep
) -> Any:
    """RFC 7591 dynamic client registration.

    Open on purpose, and harmless on its own: a registered client can do
    nothing until a signed-in human approves it on the consent page.
    """
    if not settings.MCP_ENABLED:
        return _error("access_denied", "The MCP endpoint is switched off.", 403)
    if _too_many(f"register:{_client_ip(request)}", limit=10, window=3600):
        return _error("invalid_request", "Too many registrations.", 429)

    if not 1 <= len(body.redirect_uris) <= 10:
        return _error("invalid_redirect_uri", "Give between 1 and 10 redirect URIs.")
    for uri in body.redirect_uris:
        if not _valid_redirect_uri(uri):
            return _error(
                "invalid_redirect_uri",
                f"'{uri}' must be https, or http on a loopback address, "
                "and carry no fragment.",
            )
    if body.token_endpoint_auth_method not in (None, "none"):
        return _error(
            "invalid_client_metadata",
            "Only public clients are supported; use PKCE rather than a secret.",
        )
    if body.response_types not in (None, ["code"]):
        return _error("invalid_client_metadata", "Only the 'code' response type.")
    if body.grant_types is not None and not set(body.grant_types) <= {
        "authorization_code",
        "refresh_token",
    }:
        return _error(
            "invalid_client_metadata",
            "Only the authorization_code and refresh_token grants.",
        )

    client = OAuthClient(
        client_name=body.client_name, redirect_uris=list(body.redirect_uris)
    )
    session.add(client)
    session.commit()
    session.refresh(client)

    return JSONResponse(
        status_code=201,
        content=OAuthClientInfo(
            client_id=str(client.id),
            client_name=client.client_name,
            redirect_uris=client.redirect_uris,
            client_id_issued_at=int(client.created_at.timestamp()),
        ).model_dump(),
        headers=_NO_STORE,
    )


# -----------------------------------------------------------------------------
# Authorization
# -----------------------------------------------------------------------------


def _load_client(session: SessionDep, client_id: str) -> OAuthClient | None:
    try:
        return session.get(OAuthClient, uuid.UUID(client_id))
    except ValueError:
        return None


@router.get(
    "/authorize/validate",
    response_model=OAuthAuthorizeInfo,
    # Signed in, but which user it is does not matter until they approve.
    dependencies=[Depends(get_current_user)],
)
def authorize_validate(
    client_id: str,
    redirect_uri: str,
    session: SessionDep,
) -> Any:
    """What the consent page should say, checked before it says it."""
    if not settings.MCP_ENABLED:
        return _error("access_denied", "The MCP endpoint is switched off.", 403)
    client = _load_client(session, client_id)
    if client is None:
        return _error("invalid_client", "Unknown client.")
    if redirect_uri not in client.redirect_uris:
        return _error("invalid_request", "That redirect URI is not registered.")
    return OAuthAuthorizeInfo(
        client_name=client.client_name,
        redirect_uri=redirect_uri,
        scope=security.MCP_SCOPE,
    )


@router.post("/authorize", response_model=OAuthAuthorizeResponse)
def authorize(
    body: OAuthAuthorizeRequest,
    current_user: CurrentUser,
    session: SessionDep,
) -> Any:
    """Approve a client, on behalf of the signed-in user."""
    if not settings.MCP_ENABLED:
        return _error("access_denied", "The MCP endpoint is switched off.", 403)

    client = _load_client(session, body.client_id)
    if client is None:
        return _error("invalid_client", "Unknown client.")
    if body.redirect_uri not in client.redirect_uris:
        return _error("invalid_request", "That redirect URI is not registered.")
    if body.code_challenge_method != "S256":
        return _error("invalid_request", "PKCE must use S256.")
    if not _PKCE_RE.match(body.code_challenge):
        return _error("invalid_request", "Malformed code challenge.")
    if body.resource and body.resource.rstrip("/") != settings.mcp_resource:
        return _error("invalid_target", "Unknown resource.")

    code = secrets.token_urlsafe(32)
    session.add(
        OAuthAuthorizationCode(
            code_hash=_hash(code),
            client_id=client.id,
            user_id=current_user.id,
            redirect_uri=body.redirect_uri,
            code_challenge=body.code_challenge,
            resource=body.resource,
            expires_at=_now() + timedelta(seconds=settings.OAUTH_CODE_EXPIRE_SECONDS),
        )
    )
    session.commit()

    query = {"code": code}
    if body.state:
        query["state"] = body.state
    separator = "&" if urlparse(body.redirect_uri).query else "?"
    # Built from the registered URI, never from what the request asked for.
    return OAuthAuthorizeResponse(
        redirect_url=f"{body.redirect_uri}{separator}{urlencode(query)}"
    )


# -----------------------------------------------------------------------------
# Token
# -----------------------------------------------------------------------------


@router.post("/token")
def token(  # noqa: PLR0911 - each branch is a distinct OAuth error
    request: Request,
    session: SessionDep,
    grant_type: str = Form(...),
    code: str | None = Form(None),
    redirect_uri: str | None = Form(None),
    client_id: str | None = Form(None),
    code_verifier: str | None = Form(None),
    refresh_token: str | None = Form(None),
    resource: str | None = Form(None),
) -> Any:
    """Exchange a code, or a refresh token, for an access token."""
    if not settings.MCP_ENABLED:
        return _error("access_denied", "The MCP endpoint is switched off.", 403)
    if _too_many(f"token:{_client_ip(request)}", limit=60, window=60):
        return _error("invalid_request", "Too many token requests.", 429)

    if grant_type == "authorization_code":
        return _authorization_code_grant(
            session, code, redirect_uri, client_id, code_verifier, resource
        )
    if grant_type == "refresh_token":
        return _refresh_token_grant(session, refresh_token, client_id)
    return _error("unsupported_grant_type", f"'{grant_type}' is not supported.")


def _authorization_code_grant(
    session: SessionDep,
    code: str | None,
    redirect_uri: str | None,
    client_id: str | None,
    code_verifier: str | None,
    resource: str | None,
) -> Any:
    if not (code and redirect_uri and client_id and code_verifier):
        return _error("invalid_request", "Missing a required parameter.")

    record = session.exec(
        select(OAuthAuthorizationCode).where(
            OAuthAuthorizationCode.code_hash == _hash(code)
        )
    ).first()
    if record is None:
        return _error("invalid_grant", "Unknown or already used code.")

    if record.used_at is not None:
        # OAuth 2.1 §4.1.3: a replayed code means the first exchange may have
        # been someone else's, so everything it produced is withdrawn.
        _revoke_family(session, record.refresh_token_id)
        session.commit()
        return _error("invalid_grant", "That code has already been used.")
    if _aware(record.expires_at) < _now():
        return _error("invalid_grant", "That code has expired.")
    if str(record.client_id) != client_id or record.redirect_uri != redirect_uri:
        return _error("invalid_grant", "The code was issued to another client.")
    if resource and resource.rstrip("/") != (record.resource or "").rstrip("/"):
        return _error("invalid_target", "Resource does not match the authorization.")

    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
    expected = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
    if not hmac.compare_digest(expected, record.code_challenge):
        return _error("invalid_grant", "PKCE verification failed.")

    user = session.get(User, record.user_id)
    if user is None or not user.is_active:
        return _error("invalid_grant", "That account is no longer active.")

    payload, refresh = _issue(session, user.id, record.client_id, uuid.uuid4())
    record.used_at = _now()
    record.refresh_token_id = refresh.id
    session.add(record)
    _prune(session)
    session.commit()
    return JSONResponse(content=payload, headers=_NO_STORE)


def _refresh_token_grant(
    session: SessionDep, refresh_token: str | None, client_id: str | None
) -> Any:
    if not refresh_token:
        return _error("invalid_request", "Missing the refresh token.")

    record = session.exec(
        select(OAuthRefreshToken).where(
            OAuthRefreshToken.token_hash == _hash(refresh_token)
        )
    ).first()
    if record is None:
        return _error("invalid_grant", "Unknown refresh token.")
    if record.revoked:
        # A revoked token coming back means someone kept a copy: the whole
        # family goes, including whatever is in legitimate use.
        _revoke_family(session, record.id, family_id=record.family_id)
        session.commit()
        return _error("invalid_grant", "That refresh token was already used.")
    if _aware(record.expires_at) < _now():
        return _error("invalid_grant", "That refresh token has expired.")
    if client_id and str(record.client_id) != client_id:
        return _error("invalid_grant", "The token belongs to another client.")

    user = session.get(User, record.user_id)
    if user is None or not user.is_active:
        return _error("invalid_grant", "That account is no longer active.")

    payload, _ = _issue(session, user.id, record.client_id, record.family_id)
    record.revoked = True
    session.add(record)
    session.commit()
    return JSONResponse(content=payload, headers=_NO_STORE)


def _revoke_family(
    session: SessionDep,
    token_id: uuid.UUID | None,
    family_id: uuid.UUID | None = None,
) -> None:
    """Withdraw every refresh token descended from one authorization."""
    if family_id is None:
        if token_id is None:
            return
        issued = session.get(OAuthRefreshToken, token_id)
        if issued is None:
            return
        family_id = issued.family_id
    for token in session.exec(
        select(OAuthRefreshToken).where(OAuthRefreshToken.family_id == family_id)
    ).all():
        token.revoked = True
        session.add(token)


# -----------------------------------------------------------------------------
# Management
#
# Registration is open, so without these the only way to withdraw one agent's
# access was rotating the signing key — which cuts off every agent at once.
# They answer in FastAPI's error shape rather than OAuth's, because the
# dashboard reads them and no OAuth client ever does, and they keep working
# with MCP switched off: that is exactly when leftover clients want clearing.
# -----------------------------------------------------------------------------


class RegisteredClient(BaseModel):
    """A registered agent, and whether anyone actually let it in."""

    id: uuid.UUID
    client_name: str
    redirect_uris: list[str]
    created_at: datetime
    #: Live refresh tokens. Zero means it registered and was never approved.
    active_tokens: int
    last_authorized_at: datetime | None


class RegisteredClients(BaseModel):
    data: list[RegisteredClient]
    count: int


@router.get(
    "/clients",
    response_model=RegisteredClients,
    dependencies=[Depends(get_current_active_superuser)],
)
def read_clients(session: SessionDep) -> Any:
    """Every agent that registered itself, newest first."""
    # ponytail: counts live tokens in Python; a GROUP BY if this ever grows.
    # It also keeps the expiry comparison off a naive Postgres column.
    now = _now()
    live: dict[uuid.UUID, list[datetime]] = defaultdict(list)
    for token in session.exec(
        select(OAuthRefreshToken).where(col(OAuthRefreshToken.revoked).is_(False))
    ).all():
        if _aware(token.expires_at) > now:
            live[token.client_id].append(_aware(token.created_at))

    data = [
        RegisteredClient(
            id=client.id,
            client_name=client.client_name,
            redirect_uris=client.redirect_uris,
            created_at=_aware(client.created_at),
            active_tokens=len(live[client.id]),
            last_authorized_at=max(live[client.id], default=None),
        )
        for client in session.exec(
            select(OAuthClient).order_by(col(OAuthClient.created_at).desc())
        ).all()
    ]
    return RegisteredClients(data=data, count=len(data))


@router.delete(
    "/clients/{client_id}",
    response_model=Message,
    dependencies=[Depends(get_current_active_superuser)],
)
def revoke_client(client_id: uuid.UUID, session: SessionDep) -> Any:
    """Withdraw one agent's access, leaving every other agent alone.

    Deleting the client cascades to its codes and refresh tokens, so it can
    get nothing new and cannot come back without registering again. An access
    token already in its hands keeps working until it expires
    (``MCP_TOKEN_EXPIRE_MINUTES``) — those are stateless by design.
    """
    client = session.get(OAuthClient, client_id)
    if client is None:
        raise HTTPException(status_code=404, detail="No such client")
    name = client.client_name
    session.delete(client)
    session.commit()
    return Message(message=f"Revoked '{name}'")