Logan
030dd2d787
app/internal/storage.py Replaced make_public() + public_url with a v2 signed URL (1-year expiry, no public bucket needed) app/main.py Releases all in-use tokens at startup — tokens from previous sessions are cleared automatically app/routers/tokens.py Added POST /tokens/flush to force-release orphaned tokens on demand
127 lines
4.1 KiB
Python
127 lines
4.1 KiB
Python
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,
|
|
})
|