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

"""
Per-bot OAuth "Connect a Discord server" (server-26#225 S6, owner decision Q1).

An org owner links a Discord server by inviting ONE radio bot through
Discord's OAuth2 bot invite. The invite's ``state`` carries a single-use,
short-lived nonce bound to the org and the user who started it; Discord
redirects back with ``code`` (and ``guild_id``), C2 exchanges the code with
THAT bot application's client secret, and the exchange's ``guild`` object —
never the query string — says which server to link.

Secrets:

  bot_oauth/{token_id}            server-only (Firestore rules deny all
    client_secret, enabled,       client reads; c2-core's admin SDK only)
    set_by, set_at, checked_at

  The secret is never on ``bot_tokens`` (GET /tokens returns an allowlist of
  those fields), never returned by any route, never logged, never put in an
  error string. Firestore encrypts it at rest like every other document; the
  codebase has no application-level encryption pattern (bot tokens are
  stored the same way), so none is invented here.

States:

  discord_oauth_states/{sha256(nonce)}   the nonce itself is never stored
    org_id, uid, token_id, node_id, return_to, created_at, expires_at,
    used_at|None, used_by|None

OAuth is off while ``settings.discord_oauth_redirect_uri`` is empty: every
route 404s and the UI hides the button.

HTTP goes through discord_rest's shared client (User-Agent, transport for
tests) and counts 401/403/429 answers against its process-wide breaker, so a
wrong secret retried in a loop can't contribute to an IP ban either.
"""
from __future__ import annotations

import hashlib
import secrets
import time
from datetime import datetime, timedelta, timezone
from typing import Any, Optional
from urllib.parse import urlencode

import httpx

from app.config import settings
from app.internal import discord_leases as leases
from app.internal import discord_rest
from app.internal import firestore as fstore

OAUTH = "bot_oauth"
STATES = "discord_oauth_states"
AUTHORIZE_URL = "https://discord.com/oauth2/authorize"
# Discord client secrets are 32 url-safe characters today; accept a margin.
SECRET_MIN, SECRET_MAX = 16, 128
RETURN_TO_MAX = 200


def enabled() -> bool:
    return bool((settings.discord_oauth_redirect_uri or "").strip())


def _now() -> datetime:
    return datetime.now(timezone.utc)


def state_key(nonce: str) -> str:
    return hashlib.sha256((nonce or "").encode()).hexdigest()


def safe_return_to(value: Optional[str]) -> Optional[str]:
    """A same-site path only (``/nodes/abc``), never an absolute or
    protocol-relative URL — the callback page redirects to it."""
    v = (value or "").strip()
    if not v or len(v) > RETURN_TO_MAX or not v.startswith("/") or v.startswith("//") or "\\" in v:
        return None
    if any(ord(ch) < 0x20 for ch in v):
        return None
    return v


# ---------------------------------------------------------------------------
# Secrets
# ---------------------------------------------------------------------------

async def enabled_token_ids() -> set[str]:
    """Bots whose OAuth secret is set and enabled."""
    return {d.get("token_id") for d in await fstore.collection_list(OAUTH)
            if d.get("token_id") and d.get("enabled") and d.get("client_secret")}


async def load_secret(token_id: str) -> Optional[str]:
    doc = await fstore.doc_get(OAUTH, token_id)
    if not doc or not doc.get("enabled") or not doc.get("client_secret"):
        return None
    return doc["client_secret"]


async def store_secret(token_id: str, client_secret: str, uid: str) -> None:
    now = _now().isoformat()
    await fstore.doc_set(OAUTH, token_id, {
        "token_id": token_id, "client_secret": client_secret, "enabled": True,
        "set_by": uid, "set_at": now, "checked_at": now,
    }, merge=False)


async def clear_secret(token_id: str) -> bool:
    if not await fstore.doc_get(OAUTH, token_id):
        return False
    await fstore.doc_delete(OAUTH, token_id)
    return True


# ---------------------------------------------------------------------------
# Discord OAuth2 HTTP (client_id + client_secret, basic auth)
# ---------------------------------------------------------------------------

async def _post_form(path: str, data: dict, client_id: str, client_secret: str) -> discord_rest.RestResult:
    """POST an application-authenticated form. Never raises for a Discord or
    network failure; the error string never holds the secret, the code or a
    token (only Discord's short ``error`` field)."""
    now = time.monotonic()
    if discord_rest._breaker.tripped(now):
        return discord_rest.RestResult(None, error="Discord REST paused (invalid-request budget reached)")
    try:
        r = await discord_rest._client().post(path, data=data, auth=(client_id, client_secret))
    except httpx.HTTPError as e:
        return discord_rest.RestResult(None, error=type(e).__name__)
    status = r.status_code
    if status in (401, 403) or (status == 429 and r.headers.get("x-ratelimit-scope") != "shared"):
        discord_rest._breaker.record(time.monotonic())
    body: Any = None
    try:
        body = r.json() if r.content else None
    except ValueError:
        body = None
    if 200 <= status < 300:
        return discord_rest.RestResult(status, body)
    err = body.get("error") if isinstance(body, dict) and isinstance(body.get("error"), str) else None
    return discord_rest.RestResult(status, None, error=f"Discord answered {status}" + (f" ({err[:40]})" if err else ""))


async def exchange_code(client_id: str, client_secret: str, code: str) -> discord_rest.RestResult:
    return await _post_form("/oauth2/token", {
        "grant_type": "authorization_code", "code": code,
        "redirect_uri": settings.discord_oauth_redirect_uri.strip(),
    }, client_id, client_secret)


async def revoke(client_id: str, client_secret: str, access_token: Optional[str]) -> None:
    """Best effort: C2 never uses the user's access token, so it is revoked
    as soon as the guild has been read from the exchange."""
    if access_token:
        await _post_form("/oauth2/token/revoke", {"token": access_token, "token_type_hint": "access_token"},
                         client_id, client_secret)


async def check_secret(client_id: str, client_secret: str) -> discord_rest.RestResult:
    """Prove ``client_secret`` belongs to ``client_id`` with a client-
    credentials grant (scope ``identify`` — the bot owner's own token), then
    revoke the token it returns. Nothing is stored from the answer."""
    r = await _post_form("/oauth2/token", {"grant_type": "client_credentials", "scope": "identify"},
                         client_id, client_secret)
    if r.ok and isinstance(r.data, dict):
        await revoke(client_id, client_secret, r.data.get("access_token"))
    return r


# ---------------------------------------------------------------------------
# State + authorize URL
# ---------------------------------------------------------------------------

def authorize_url(application_id: str, nonce: str) -> str:
    return AUTHORIZE_URL + "?" + urlencode({
        "client_id": application_id,
        "scope": "bot",
        "permissions": str(leases.RADIO_BOT_PERMISSIONS),
        "response_type": "code",
        "redirect_uri": settings.discord_oauth_redirect_uri.strip(),
        "state": nonce,
    })


async def create_state(*, org_id: str, uid: str, token_id: str, node_id: Optional[str],
                       return_to: Optional[str]) -> tuple[str, datetime]:
    nonce = secrets.token_urlsafe(32)
    now = _now()
    expires_at = now + timedelta(seconds=max(60, settings.discord_oauth_state_ttl_seconds))
    await fstore.doc_set(STATES, state_key(nonce), {
        "org_id": org_id, "uid": uid, "token_id": token_id, "node_id": node_id,
        "return_to": return_to, "created_at": now.isoformat(), "expires_at": expires_at.isoformat(),
        "used_at": None, "used_by": None,
    }, merge=False)
    return nonce, expires_at


class StateError(Exception):
    def __init__(self, status: int, code: str, message: str):
        super().__init__(message)
        self.status, self.code, self.message = status, code, message


def consume_state_txn(nonce: str, uid: str, now: datetime):
    """Transaction body: the state must exist, be unexpired and unused, and
    belong to ``uid`` (login-CSRF: someone else's state never links into the
    caller's org, and the caller's own code can't be replayed by another
    account). Any presentation of a live state burns it — a failure later in
    the flow means starting again, never retrying the same state."""
    db = fstore.db
    key = state_key(nonce)

    def _txn(txn):
        ref = db.collection(STATES).document(key)
        snap = ref.get(transaction=txn)
        doc = snap.to_dict() if snap.exists else None
        if not doc:
            return {"error": StateError(400, "invalid_state",
                                        "That Discord link is not valid. Start again from DRB.")}
        if doc.get("used_at"):
            return {"error": StateError(409, "state_used",
                                        "That Discord link was already used. Start again from DRB.")}
        txn.update(ref, {"used_at": now.isoformat(), "used_by": uid})
        try:
            expired = now > datetime.fromisoformat(doc.get("expires_at") or "")
        except ValueError:
            expired = True
        if expired:
            return {"error": StateError(410, "state_expired",
                                        "That Discord link expired. Start again from DRB.")}
        if doc.get("uid") != uid:
            return {"error": StateError(403, "wrong_user",
                                        "That Discord link was started by a different DRB account. "
                                        "Sign in as that account, or start again.")}
        return {"state": doc}

    return _txn


async def consume_state(nonce: str, uid: str) -> dict:
    if not nonce or len(nonce) > 200:
        raise StateError(400, "invalid_state", "That Discord link is not valid. Start again from DRB.")
    out = await fstore.run_transaction(consume_state_txn(nonce, uid, _now()))
    if "error" in out:
        raise out["error"]
    return out["state"]