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

import asyncio
import random
import string
from datetime import datetime, timezone, timedelta
from uuid import uuid4
from fastapi import APIRouter, HTTPException, Depends, Request
from firebase_admin import auth as firebase_auth
from pydantic import BaseModel
from app.internal import firestore as fstore
from app.internal.auth import require_firebase_token, require_service_key, get_role
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}


# ---------------------------------------------------------------------------
# Org provisioning — SAAS_PLAN.md B4. The client creates the Firebase user
# first (email/password or Google) and calls this with that user's fresh ID
# token, which carries no org_id/org_role claim yet. This is the only route
# that turns "has a Firebase account" into "can read anything" — see
# infra/firestore/firestore.rules and AuthProvider's no-claim guard.
# ---------------------------------------------------------------------------

class SignupBody(BaseModel):
    org_name: str


@router.post("/signup")
async def signup(body: SignupBody, decoded: dict = Depends(require_firebase_token)):
    """
    Provision a new organization owned by the calling user, or return their
    existing one. Idempotent by design: the frontend calls this right after
    account creation, and a user who double-submits (or re-runs it after a
    refresh) must not end up with two orgs.
    """
    uid = decoded["uid"]

    existing_org_id = decoded.get("org_id")
    if existing_org_id:
        org = await fstore.doc_get("organizations", existing_org_id)
        if org:
            return {"org_id": existing_org_id, "org_name": org.get("name"), "already_provisioned": True}
        # Claim points at a deleted/missing org doc — fall through and
        # provision a fresh one rather than leaving the account stranded.

    org_name = body.org_name.strip()
    if not org_name:
        raise HTTPException(400, "org_name is required.")
    if len(org_name) > 200:
        raise HTTPException(400, "org_name is too long.")

    org_id = str(uuid4())
    now = datetime.now(timezone.utc).isoformat()

    # plan_id/subscription_status/stripe_*/seat_limit/node_limit/retention_days
    # are all deliberately None — no billing model exists yet (see
    # app/internal/tenancy.py). This is the seam a future billing pass writes
    # into; nothing today reads or enforces these fields.
    await fstore.doc_set("organizations", org_id, {
        "org_id": org_id,
        "name": org_name,
        "created_at": now,
        "created_by_uid": uid,
        "plan_id": None,
        "subscription_status": None,
        "stripe_customer_id": None,
        "stripe_subscription_id": None,
        "current_period_end": None,
        "seat_limit": None,
        "node_limit": None,
        "retention_days": None,
    }, merge=False)

    await fstore.doc_set("org_members", uid, {
        "uid": uid,
        "org_id": org_id,
        "org_role": "owner",
        "email": decoded.get("email"),
        "added_at": now,
    }, merge=False)

    # set_custom_user_claims() replaces the whole claim set, so preserve any
    # existing custom claims (owned_node_ids, a platform `role` if this
    # account was created via the admin-only POST /admin/users flow, etc.)
    # rather than clobbering them. Firebase's own reserved JWT fields are
    # stripped out — they aren't settable as custom claims and would raise.
    _RESERVED = {
        "iss", "aud", "auth_time", "user_id", "sub", "iat", "exp", "uid",
        "email", "email_verified", "firebase", "name", "picture",
    }
    existing_claims = {k: v for k, v in decoded.items() if k not in _RESERVED}
    # role: platform-level, orthogonal to org ownership. get_role() falls
    # back to "viewer" for a brand-new self-serve signup with no claims yet.
    claims = {**existing_claims, "org_id": org_id, "org_role": "owner", "role": get_role(decoded)}
    await asyncio.to_thread(firebase_auth.set_custom_user_claims, uid, claims)

    logger.info(f"Org provisioned: org_id={org_id} name={org_name!r} owner_uid={uid}")
    return {"org_id": org_id, "org_name": org_name, "already_provisioned": False}


# ---------------------------------------------------------------------------
# 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}