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, })