import uuid from fastapi import APIRouter, HTTPException from pydantic import BaseModel from typing import Optional from datetime import datetime, timezone from app.internal import firestore as fstore 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"][:10] + "…" + t["token"][-4:]} for t in tokens ] @router.post("", status_code=201) async def add_token(body: TokenCreate): 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(): """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.delete("/{token_id}", status_code=204) async def delete_token(token_id: str): 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, })