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

import random
import string
from datetime import datetime, timezone, timedelta
from uuid import uuid4
from fastapi import APIRouter, HTTPException, Depends, Request
from pydantic import BaseModel
from app.internal import firestore as fstore
from app.internal.auth import require_firebase_token, require_service_key
from app.internal.logger import logger

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

_CODE_TTL_MINUTES = 15


def _gen_code() -> str:
    return "".join(random.choices(string.ascii_uppercase + string.digits, k=6))


# ---------------------------------------------------------------------------
# Web: generate a short-lived linking code
# ---------------------------------------------------------------------------

@router.post("/link/generate")
async def generate_link_code(decoded: dict = Depends(require_firebase_token)):
    """Authenticated Firebase user generates a code to paste into Discord /link."""
    firebase_uid = decoded["uid"]

    # Check if already linked
    existing = await fstore.doc_get("firebase_discord_links", firebase_uid)
    if existing and existing.get("discord_user_id"):
        return {
            "already_linked": True,
            "discord_user_id": existing["discord_user_id"],
        }

    code = _gen_code()
    expires_at = (datetime.now(timezone.utc) + timedelta(minutes=_CODE_TTL_MINUTES)).isoformat()
    await fstore.doc_set("link_codes", code, {
        "firebase_uid": firebase_uid,
        "expires_at": expires_at,
    }, merge=False)

    return {"code": code, "expires_minutes": _CODE_TTL_MINUTES}


# ---------------------------------------------------------------------------
# Discord bot: resolve a code and store the link
# ---------------------------------------------------------------------------

class LinkResolveBody(BaseModel):
    code: str
    discord_user_id: str
    discord_username: str = ""


@router.post("/link")
async def resolve_link_code(body: LinkResolveBody, _: dict = Depends(require_service_key)):
    """Discord bot resolves a linking code and permanently links the accounts."""
    doc = await fstore.doc_get("link_codes", body.code.upper().strip())
    if not doc:
        raise HTTPException(404, "Invalid or expired code.")

    expires_at = datetime.fromisoformat(doc["expires_at"])
    if datetime.now(timezone.utc) > expires_at:
        await fstore.doc_delete("link_codes", body.code)
        raise HTTPException(410, "Code has expired. Generate a new one from the web app.")

    firebase_uid = doc["firebase_uid"]

    # Check if this Discord account is already linked to a different Firebase UID
    existing = await fstore.doc_get("discord_links", body.discord_user_id)
    if existing and existing.get("firebase_uid") and existing["firebase_uid"] != firebase_uid:
        raise HTTPException(409, "This Discord account is already linked to a different account.")

    now = datetime.now(timezone.utc).isoformat()

    # Store both directions
    await fstore.doc_set("discord_links", body.discord_user_id, {
        "firebase_uid": firebase_uid,
        "discord_username": body.discord_username,
        "linked_at": now,
    }, merge=False)

    await fstore.doc_set("firebase_discord_links", firebase_uid, {
        "discord_user_id": body.discord_user_id,
        "discord_username": body.discord_username,
        "linked_at": now,
    }, merge=False)

    # Clean up the code
    await fstore.doc_delete("link_codes", body.code)

    logger.info(f"Linked firebase_uid={firebase_uid} <-> discord_user_id={body.discord_user_id}")
    return {"ok": True, "firebase_uid": firebase_uid}


# ---------------------------------------------------------------------------
# Web: check current link status
# ---------------------------------------------------------------------------

@router.get("/link/status")
async def link_status(decoded: dict = Depends(require_firebase_token)):
    firebase_uid = decoded["uid"]
    link = await fstore.doc_get("firebase_discord_links", firebase_uid)
    if link and link.get("discord_user_id"):
        return {
            "linked": True,
            "discord_user_id": link["discord_user_id"],
            "discord_username": link.get("discord_username", ""),
            "linked_at": link.get("linked_at"),
        }
    return {"linked": False}


# ---------------------------------------------------------------------------
# Web: unlink
# ---------------------------------------------------------------------------

@router.delete("/link")
async def unlink(decoded: dict = Depends(require_firebase_token)):
    firebase_uid = decoded["uid"]
    link = await fstore.doc_get("firebase_discord_links", firebase_uid)
    if not link or not link.get("discord_user_id"):
        raise HTTPException(404, "No linked Discord account.")
    discord_user_id = link["discord_user_id"]
    await fstore.doc_delete("discord_links", discord_user_id)
    await fstore.doc_delete("firebase_discord_links", firebase_uid)
    return {"ok": True}


# ---------------------------------------------------------------------------
# Session recording — called by the frontend on each successful sign-in
# ---------------------------------------------------------------------------

@router.post("/session")
async def record_session(request: Request, decoded: dict = Depends(require_firebase_token)):
    """Record a sign-in event for the authenticated user."""
    session_id = str(uuid4())
    ip = request.client.host if request.client else None
    user_agent = request.headers.get("user-agent", "")

    await fstore.doc_set("user_sessions", session_id, {
        "session_id": session_id,
        "uid": decoded["uid"],
        "email": decoded.get("email", ""),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "ip": ip,
        "user_agent": user_agent,
    }, merge=False)

    return {"ok": True}