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 datetime import datetime, timezone
from app.internal import firestore as fstore
from app.internal.auth import require_admin_token

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


# ---------------------------------------------------------------------------
# CRUD
# ---------------------------------------------------------------------------

@router.get("")
async def list_tokens():
    """List all tokens. The actual token string is masked for safety."""
    tokens = await fstore.collection_list("bot_tokens")
    return [
        {**t, "token": "•••" + t["token"][-4:]}
        for t in tokens
    ]


@router.post("", status_code=201)
async def add_token(body: TokenCreate, _: dict = Depends(require_admin_token)):
    token_id = str(uuid.uuid4())
    doc = {
        "token_id": token_id,
        "name": body.name,
        "token": body.token,
        "in_use": False,
        "assigned_node_id": None,
        "assigned_at": None,
    }
    await fstore.doc_set("bot_tokens", token_id, doc, merge=False)
    return {"token_id": token_id, "name": body.name}


@router.post("/flush", status_code=200)
async def flush_tokens(_: dict = Depends(require_admin_token)):
    """Force-release all in-use tokens (admin utility — use when tokens get orphaned)."""
    def _find():
        from app.internal.firestore import db
        return [d for d in db.collection("bot_tokens").where("in_use", "==", True).stream()]

    import asyncio
    results = await asyncio.to_thread(_find)
    for doc in results:
        await fstore.doc_update("bot_tokens", doc.id, {
            "in_use": False,
            "assigned_node_id": None,
            "assigned_at": None,
        })
    return {"released": len(results)}


@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("in_use"):
        raise HTTPException(409, "Token is currently in use by a node.")
    await fstore.doc_delete("bot_tokens", token_id)


# ---------------------------------------------------------------------------
# Internal helpers — used by the nodes router, not exposed via HTTP
# ---------------------------------------------------------------------------

async def assign_token(node_id: str, preferred_token_id: Optional[str] = None) -> Optional[str]:
    """
    Find a free token, mark it as in-use, return the token string.
    If preferred_token_id is given, try that token first (only if it's free).
    Returns None if no tokens are available.
    """
    def _find_free(preferred: Optional[str]):
        from app.internal.firestore import db
        if preferred:
            snap = db.collection("bot_tokens").document(preferred).get()
            if snap.exists and not snap.to_dict().get("in_use"):
                return [snap]
        docs = db.collection("bot_tokens").where("in_use", "==", False).limit(1).stream()
        return [d for d in docs]

    import asyncio
    results = await asyncio.to_thread(_find_free, preferred_token_id)
    if not results:
        return None

    doc = results[0]
    token_id = doc.id
    token_value = doc.to_dict()["token"]

    await fstore.doc_update("bot_tokens", token_id, {
        "in_use": True,
        "assigned_node_id": node_id,
        "assigned_at": datetime.now(timezone.utc),
    })
    return token_value


async def release_token(node_id: str) -> None:
    """Free whichever token is currently assigned to this node."""
    def _find_assigned():
        from app.internal.firestore import db
        return [
            d for d in db.collection("bot_tokens")
            .where("assigned_node_id", "==", node_id)
            .stream()
        ]

    import asyncio
    results = await asyncio.to_thread(_find_assigned)
    for doc in results:
        await fstore.doc_update("bot_tokens", doc.id, {
            "in_use": False,
            "assigned_node_id": None,
            "assigned_at": None,
        })