server-26/drb-c2-core/app/routers/discord_oauth.py

"""
"Connect a Discord server" through a radio bot's own OAuth2 invite
(server-26#225 S6, owner decision Q1 = per-bot OAuth).

  GET  /discord/oauth/options   org owner / platform admin — which bot Connect
                                would invite (no state minted). 404 = off.
  POST /discord/oauth/start     org owner / platform admin — mint a state,
                                return Discord's authorize URL
  POST /discord/oauth/complete  signed-in user who started it — exchange the
                                code, link the server, check permissions

Platform admins may pass ?org_id= to act for another org (resolve_org_scope).
Every route is a 404 while ``discord_oauth_redirect_uri`` is empty or no
eligible bot has an OAuth secret (PUT /tokens/{id}/oauth).

Which bot (``_pick_bot``): an explicit ``token_id`` the org may use; else the
radio bot preferred for the node's system (``node_id``); else the org's own
bring-your-own bot; else the pool bot with the fewest servers that is under
``discord_pool_bot_guild_cap`` (owner decision Q3). Only bots with an OAuth
secret set count. A bot the caller's org can't use is a 404, never a 403.

The org a server links to is the one the STATE names, re-checked against a
fresh read of the caller's membership; the server is the one Discord's code
exchange names. The redirect's ``guild_id`` query parameter is ignored.

This reverses ``discord_web``'s old "C2 never calls Discord on the browser's
behalf" on purpose (plan finding 5): the call is made with the bot
application's credentials, after the org-scope checks and rate limits here.
"""
from __future__ import annotations

import asyncio
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from app.config import settings
from app.internal import discord_guild_link as glink
from app.internal import discord_leases as leases
from app.internal import discord_oauth as oauth
from app.internal import discord_rest
from app.internal import discord_streams as streams
from app.internal import firestore as fstore
from app.internal.audit import write_audit
from app.internal.auth import (
    _RateLimiter, describe_actor, get_role, require_firebase_token, require_org_owner_token, resolve_org_scope,
)
from app.internal.logger import logger
from app.routers.discord_checks import voice_check

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

start_limiter = _RateLimiter(max_calls=5, window_seconds=600)
complete_limiter = _RateLimiter(max_calls=10, window_seconds=600)
_ID = r"^[A-Za-z0-9_-]{1,64}$"
_OFF = "Connecting a Discord server through DRB isn't available."


def _require_enabled() -> None:
    if not oauth.enabled():
        raise HTTPException(404, _OFF)


def _err(status: int, code: str, message: str, **extra) -> JSONResponse:
    return JSONResponse(status_code=status, content={"detail": message, "code": code, **extra})


def _server_count(tok: dict) -> Optional[int]:
    members = tok.get("guild_ids")
    return len(members) if isinstance(members, list) else None


def _usable(tok: dict, org_id: str, oauth_ids: set) -> bool:
    return (tok.get("token_id") in oauth_ids
            and tok.get("valid") is not False
            and bool(tok.get("application_id"))
            and tok.get("org_id") in (None, org_id)
            and leases.invitable(tok))


async def _pick_bot(org_id: str, *, node_id: Optional[str] = None,
                    token_id: Optional[str] = None) -> tuple[dict, str]:
    """(token doc, why) — why is ``chosen`` | ``preferred`` | ``org_bot`` |
    ``pool``. Raises HTTPException 404 when nothing fits."""
    tokens = {t["token_id"]: t for t in await fstore.collection_list(leases.TOKENS) if t.get("token_id")}
    oauth_ids = await oauth.enabled_token_ids()

    if token_id:
        tok = tokens.get(token_id)
        if not tok or tok.get("org_id") not in (None, org_id):
            raise HTTPException(404, "Radio bot not found.")
        if not _usable(tok, org_id, oauth_ids):
            raise HTTPException(404, f"{tok.get('name') or 'That radio bot'} can't be used to connect a server "
                                     "right now.")
        return tok, "chosen"

    if node_id:
        node = await fstore.doc_get("nodes", node_id)
        if not node or node.get("org_id") != org_id:
            raise HTTPException(404, "Node not found.")
        system_id = streams._effective_system(node)
        system = await fstore.doc_get("systems", system_id) if system_id else None
        pref = tokens.get((system or {}).get("preferred_token_id") or "")
        if pref and _usable(pref, org_id, oauth_ids):
            return pref, "preferred"

    own = sorted((t for t in tokens.values() if t.get("org_id") == org_id and _usable(t, org_id, oauth_ids)),
                 key=lambda t: (str(t.get("name") or "").lower(), t["token_id"]))
    if own:
        return own[0], "org_bot"

    cap = settings.discord_pool_bot_guild_cap
    pool = sorted((t for t in tokens.values() if t.get("org_id") is None and _usable(t, org_id, oauth_ids)),
                  key=lambda t: (_server_count(t) if _server_count(t) is not None else cap,
                                 str(t.get("name") or "").lower(), t["token_id"]))
    if pool:
        return pool[0], "pool"
    raise HTTPException(404, _OFF)


def _bot_summary(tok: dict, why: str) -> dict:
    return {"name": tok.get("name") or "DRB radio bot", "reason": why,
            "own_bot": tok.get("org_id") is not None}


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@router.get("/options")
async def oauth_options(node_id: Optional[str] = Query(None, pattern=_ID),
                        token_id: Optional[str] = Query(None, pattern=_ID),
                        org_id: Optional[str] = None,
                        decoded: dict = Depends(require_org_owner_token)):
    """Whether Connect is offered here, and which bot it would invite."""
    _require_enabled()
    scope = resolve_org_scope(decoded, org_id)
    tok, why = await _pick_bot(scope, node_id=node_id, token_id=token_id)
    return {"available": True, "bot": _bot_summary(tok, why)}


class StartBody(BaseModel):
    node_id: Optional[str] = Field(None, pattern=_ID)
    token_id: Optional[str] = Field(None, pattern=_ID)
    return_to: Optional[str] = None


@router.post("/start")
async def oauth_start(body: StartBody, org_id: Optional[str] = None,
                      decoded: dict = Depends(require_org_owner_token)):
    _require_enabled()
    scope = resolve_org_scope(decoded, org_id)
    uid, email = describe_actor(decoded)
    start_limiter.check(f"{uid}:discord_oauth_start")
    tok, why = await _pick_bot(scope, node_id=body.node_id, token_id=body.token_id)
    return_to = oauth.safe_return_to(body.return_to)
    nonce, expires_at = await oauth.create_state(org_id=scope, uid=uid, token_id=tok["token_id"],
                                                 node_id=body.node_id, return_to=return_to)
    await write_audit(uid, email, "discord.oauth.start", details={
        "org_id": scope, "token_id": tok["token_id"], "reason": why, "node_id": body.node_id})
    return {"authorize_url": oauth.authorize_url(str(tok["application_id"]), nonce),
            "expires_at": expires_at.isoformat(), "bot": _bot_summary(tok, why)}


class CompleteBody(BaseModel):
    # Lengths checked in the handler: a pydantic 422 would echo the code back.
    code: str
    state: str


async def _still_owner(decoded: dict, uid: str, org_id: str) -> bool:
    """Fresh read: the claims in an ID token can be up to an hour old."""
    if get_role(decoded) == "admin":
        return True
    member = await fstore.doc_get("org_members", uid)
    return bool(member and member.get("org_id") == org_id and member.get("org_role") == "owner")


async def _channel_summary(guild_id: str, org_id: str, node_id: Optional[str]) -> Optional[dict]:
    """Permissions right after the link, within discord_oauth_link_check_seconds.
    None = not answered in time (the picker's badges will say)."""
    link = await fstore.doc_get(streams.GUILDS, guild_id) or {}
    try:
        out = await asyncio.wait_for(
            voice_check(guild_id=guild_id, org_id=org_id, node_id=node_id,
                        channels=link.get("voice_channels") or [], channel_id=None, refresh=True,
                        with_token_ids=False),
            timeout=settings.discord_oauth_link_check_seconds)
    except Exception as e:   # timeout or anything else: the link already stands
        logger.info(f"Discord OAuth link {guild_id}: permission check skipped ({type(e).__name__})")
        return None
    channels = [{"channel_id": c["channel_id"], "name": c.get("name"), "state": c["state"],
                 "fix": next((b.get("fix") for b in c.get("bots") or [] if b.get("fix")), None)}
                for c in out.get("channels") or []]
    return {
        "enabled": out.get("enabled"),
        "channels_ready": sum(1 for c in channels if c["state"] == "ready"),
        "channels_blocked": sum(1 for c in channels if c["state"] == "needs_fix"),
        "channels": channels,
    }


@router.post("/complete")
async def oauth_complete(body: CompleteBody, decoded: dict = Depends(require_firebase_token)):
    _require_enabled()
    uid, email = describe_actor(decoded)
    complete_limiter.check(f"{uid}:discord_oauth_complete")
    if not (0 < len(body.code) <= 200):
        return _err(400, "code_rejected", "Discord didn't accept that sign-in. Start again from DRB.")

    # 1. Consume the state (single use; must be this caller's).
    try:
        state = await oauth.consume_state(body.state, uid)
    except oauth.StateError as e:
        await write_audit(uid, email, "discord.oauth.refused", details={"reason": e.code})
        return _err(e.status, e.code, e.message)
    org_id, token_id = state["org_id"], state["token_id"]
    return_to = state.get("return_to")

    # 2. Still the owner of that org (fresh read, not the token's claims).
    if not await _still_owner(decoded, uid, org_id):
        await write_audit(uid, email, "discord.oauth.refused", details={"org_id": org_id, "reason": "not_owner"})
        return _err(403, "not_owner", "Only the organization owner can connect a Discord server.")

    # 3. Exchange the code with that bot application's own secret.
    tok = await fstore.doc_get(leases.TOKENS, token_id)
    secret = await oauth.load_secret(token_id)
    if (not tok or not tok.get("application_id") or not secret
            or tok.get("org_id") not in (None, org_id) or tok.get("valid") is False):
        return _err(409, "bot_unavailable", "That radio bot can no longer connect servers. Start again from DRB.")
    client_id = str(tok["application_id"])
    r = await oauth.exchange_code(client_id, secret, body.code)
    if not r.ok or not isinstance(r.data, dict):
        await write_audit(uid, email, "discord.oauth.refused", details={
            "org_id": org_id, "token_id": token_id, "reason": "exchange_failed", "status": r.status})
        if r.status is not None and 400 <= r.status < 500 and r.status != 429:
            return _err(400, "code_rejected", "Discord didn't accept that sign-in. Start again from DRB.")
        return _err(502, "discord_unavailable", "Couldn't reach Discord to finish connecting. Start again from DRB.")
    await oauth.revoke(client_id, secret, r.data.get("access_token"))

    # 4. The server is the one Discord's exchange names — never the query string.
    guild = r.data.get("guild") if isinstance(r.data.get("guild"), dict) else None
    guild_id = str((guild or {}).get("id") or "")
    if not guild_id.isdigit():
        await write_audit(uid, email, "discord.oauth.refused", details={
            "org_id": org_id, "token_id": token_id, "reason": "no_guild"})
        return _err(400, "no_guild", "Discord didn't say which server the bot was added to. Start again from DRB.")
    guild_name = guild.get("name") if isinstance(guild.get("name"), str) else None

    # 5. Link + settle (a server another org linked is never moved: 409).
    try:
        result = await glink.link_and_settle(org_id, guild_id, guild_name, uid, "oauth",
                                             extra={"linked_with_token_id": token_id})
    except glink.LinkError as e:
        await write_audit(uid, email, "discord.guild.link_refused", details={
            "guild_id": guild_id, "org_id": org_id, "reason": e.code, "via": "oauth", "token_id": token_id})
        return _err(e.status, e.code, e.message, guild_name=guild_name, bot_name=tok.get("name"))

    # 6. The bot is in the server now: record it, drop stale cached answers, audit.
    members = tok.get("guild_ids")
    if isinstance(members, list) and guild_id not in members:
        await fstore.doc_update(leases.TOKENS, token_id, {"guild_ids": [*members, guild_id]})
    discord_rest.invalidate(guild_id=guild_id)
    await write_audit(uid, email, "discord.guild.link", details={
        "guild_id": guild_id, "org_id": org_id, "already_linked": result["already_linked"],
        "previous_org_id": result["previous_org_id"], "streams_stopped": len(result["streams_stopped"]),
        "via": "oauth", "token_id": token_id,
    })
    logger.info(f"Discord guild {guild_id} linked to org={org_id!r} via OAuth (bot {token_id}) by uid={uid}")

    # 7. Which voice channels are ready, in the same request (capped).
    check = await _channel_summary(guild_id, org_id, state.get("node_id"))
    return {
        "ok": True,
        "guild_id": guild_id,
        "guild_name": guild_name,
        "org_id": org_id,
        "bot_name": tok.get("name"),
        "already_linked": result["already_linked"],
        "streams_stopped": len(result["streams_stopped"]),
        "node_id": state.get("node_id"),
        "return_to": return_to,
        "checked": check is not None,
        "channels_ready": check["channels_ready"] if check else None,
        "channels_blocked": check["channels_blocked"] if check else None,
        "channels": check["channels"] if check else [],
    }