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

import uuid
from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel
from typing import Optional
from app.internal import firestore as fstore
from app.internal import discord_leases as leases
from app.internal import discord_rest
from app.internal.auth import require_admin_token
from app.internal.logger import logger

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


class TokenCreate(BaseModel):
    name: str        # friendly label e.g. "DRB Bot 1"
    token: str       # the actual Discord bot token
    # None = platform pool. An org id = that org's bring-your-own bot, never
    # leased to anyone else (discord_leases.pick_token).
    org_id: Optional[str] = None


def _masked(t: dict) -> dict:
    out = {**t, "token": "•••" + (t.get("token") or "")[-4:]}
    out["invite_url"] = leases.invite_url(t.get("application_id"))
    return out


async def _republish_after_release(lease: Optional[dict]) -> None:
    if lease and lease.get("node_id"):
        from app.internal.discord_streams import publish_desired
        await publish_desired(lease["node_id"])


# ---------------------------------------------------------------------------
# CRUD — platform admin only (server-26#204 item 6: the list used to be open
# to any signed-in user of any org)
# ---------------------------------------------------------------------------

@router.get("")
async def list_tokens(_: dict = Depends(require_admin_token)):
    """List all tokens. The actual token string is masked for safety."""
    return [_masked(t) for t in await fstore.collection_list("bot_tokens")]


@router.post("", status_code=201)
async def add_token(body: TokenCreate, _: dict = Depends(require_admin_token)):
    """Add a radio bot. The token is checked against Discord first
    (server-26#204 item 5): a string Discord answers 401 for is refused; if
    Discord is merely unreachable the token is stored with membership unknown."""
    try:
        info = await discord_rest.inspect_token(body.token.strip())
    except discord_rest.TokenRejected as e:
        raise HTTPException(400, str(e))
    token_id = str(uuid.uuid4())
    doc = {
        "token_id": token_id,
        "name": body.name,
        "token": body.token.strip(),
        "org_id": body.org_id,
        "lease": None,
        "in_use": False,
        "assigned_node_id": None,
        "assigned_at": None,
        **info,
    }
    await fstore.doc_set("bot_tokens", token_id, doc, merge=False)
    return {"token_id": token_id, "name": body.name, "valid": info.get("valid"),
            "guild_ids": info.get("guild_ids"), "check_error": info.get("check_error"),
            "invite_url": leases.invite_url(info.get("application_id"))}


@router.post("/{token_id}/refresh", status_code=200)
async def refresh_token(token_id: str, _: dict = Depends(require_admin_token)):
    """Re-check a bot with Discord (valid? which servers is it in?)."""
    if not await fstore.doc_get("bot_tokens", token_id):
        raise HTTPException(404, "Token not found.")
    await discord_rest.refresh_membership([token_id])
    return _masked(await fstore.doc_get("bot_tokens", token_id) or {})


@router.post("/flush", status_code=200)
async def flush_tokens(_: dict = Depends(require_admin_token)):
    """Force-release every leased token (admin last resort). Each holding node
    is sent a desired-state list without the token, so it logs the bot out;
    persistent streams that lost their token are re-placed by the reconciler."""
    released = 0
    for t in await fstore.collection_list("bot_tokens"):
        if not (t.get("lease") or t.get("in_use")):
            continue
        lease = await leases.force_release_token(t["token_id"])
        released += 1
        await _republish_after_release(lease)
    logger.warning(f"Admin flushed {released} radio-bot token lease(s)")
    return {"released": released}


@router.put("/{token_id}/prefer/{system_id}", status_code=200)
async def set_preferred_system(
    token_id: str,
    system_id: str,
    _: dict = Depends(require_admin_token),
):
    """
    Mark this token as the preferred bot for a system.
    When a discord_join is issued for any node in that system, this token
    is tried first before falling back to the general pool.
    Pass system_id="_none" to clear the preference.
    """
    existing = await fstore.doc_get("bot_tokens", token_id)
    if not existing:
        raise HTTPException(404, "Token not found.")
    if system_id == "_none":
        # Clear any existing preference on the system that pointed to this token.
        system_doc = await fstore.doc_get("systems", existing.get("preferred_for_system_id", ""))
        if system_doc:
            await fstore.doc_set("systems", existing["preferred_for_system_id"], {"preferred_token_id": None})
        await fstore.doc_set("bot_tokens", token_id, {"preferred_for_system_id": None})
        return {"ok": True, "preferred_for_system_id": None}

    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        raise HTTPException(404, "System not found.")
    # Set preference on both sides for easy lookup in either direction.
    await fstore.doc_set("systems", system_id, {"preferred_token_id": token_id})
    await fstore.doc_set("bot_tokens", token_id, {"preferred_for_system_id": system_id})
    return {"ok": True, "preferred_for_system_id": system_id}


@router.delete("/{token_id}", status_code=204)
async def delete_token(token_id: str, _: dict = Depends(require_admin_token)):
    existing = await fstore.doc_get("bot_tokens", token_id)
    if not existing:
        raise HTTPException(404, "Token not found.")
    if existing.get("lease"):
        raise HTTPException(409, "Token is currently leased to a node. Stop its radio channels (or flush) first.")
    await fstore.doc_delete("bot_tokens", token_id)