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

"""
Node self-enrollment — the public replacement for the old shared-MQTT-
password flow (see MQTT-PUBLIC-AUTH-PLAN.md "Enrollment flow").

  1. POST /nodes/enroll          (X-Enrollment-Token: <fleet-wide token>)
     First-boot node upserts itself as `approval_status: pending` and gets
     back a one-time pickup_secret. Only its hash is persisted.
  2. GET  /nodes/{id}/credentials (X-Pickup-Secret: <secret from step 1>)
     Node polls this with backoff until an admin approves it in the
     frontend (existing nodes.py approve_node() flow — unchanged, still
     writes node_keys/{id}.api_key) and then reads its api_key back.

These two endpoints are meant to be public (unlike the dynsec control-plane
traffic in app/internal/dynsec.py, which never leaves the docker-internal
MQTT bridge) — that's the whole point of moving off WireGuard-per-node.
Auth is the token headers checked inline below, not the app-wide
Firebase/service-key dependency the rest of routers/nodes.py uses.
"""
import hashlib
import secrets
import time
from typing import Optional
from fastapi import APIRouter, HTTPException, Header, Request
from pydantic import BaseModel
from app.config import settings
from app.internal import firestore as fstore
from app.internal.logger import logger
from app.internal.tenancy import FOUNDING_ORG_ID

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


# ---------------------------------------------------------------------------
# Per-source-IP token bucket for /nodes/enroll.
#
# c2-core has no rate-limiting dependency anywhere today (see
# MQTT-PUBLIC-AUTH-PLAN.md); this is a deliberately small (~30 line)
# in-memory limiter rather than a new library. Known limitations:
#   - per-process: with more than one c2-core instance, each has its own
#     bucket, so real throughput is (limit x instance count). Fine today —
#     there is exactly one instance.
#   - resets on every restart/redeploy — not persisted anywhere.
# Good enough to blunt casual guessing of node_ids against the fleet token;
# not a substitute for a real edge/WAF rate limiter if this endpoint is
# ever seriously targeted.
# ---------------------------------------------------------------------------
class _TokenBucket:
    def __init__(self, capacity: int, refill_per_sec: float):
        self.capacity = capacity
        self.refill_per_sec = refill_per_sec
        self._buckets: dict[str, tuple[float, float]] = {}  # key -> (tokens, last_refill_ts)

    def allow(self, key: str) -> bool:
        now = time.monotonic()
        tokens, last_ts = self._buckets.get(key, (float(self.capacity), now))
        tokens = min(self.capacity, tokens + (now - last_ts) * self.refill_per_sec)
        if tokens < 1:
            self._buckets[key] = (tokens, now)
            return False
        self._buckets[key] = (tokens - 1, now)
        return True


# Burst of 5, refilling 1/minute — enrollment is a first-boot, once-per-node
# event, so a legitimate node never needs more than a handful of attempts.
_enroll_limiter = _TokenBucket(capacity=5, refill_per_sec=1 / 60)


def _hash_secret(secret: str) -> str:
    return hashlib.sha256(secret.encode()).hexdigest()


async def _resolve_org_for_token(token: str) -> Optional[str]:
    """
    Resolve an X-Enrollment-Token to the org_id it enrolls a node into.

    Tries the per-org enrollment_tokens collection first (SAAS_PLAN.md B2b —
    minted/listed/revoked via routers/org.py), then falls back to the legacy
    fleet-wide settings.enrollment_token so an already-deployed field node's
    .env doesn't start failing the day per-org tokens ship. The fallback
    always resolves to FOUNDING_ORG_ID — see app/internal/tenancy.py.
    """
    token_hash = _hash_secret(token)
    doc = await fstore.doc_get("enrollment_tokens", token_hash)
    if doc and not doc.get("revoked"):
        try:
            await fstore.doc_update("enrollment_tokens", token_hash, {"uses": (doc.get("uses") or 0) + 1})
        except Exception:
            pass  # use-counter is informational only — never block enrollment on it
        return doc.get("org_id")

    if settings.enrollment_token and secrets.compare_digest(token, settings.enrollment_token):
        return FOUNDING_ORG_ID

    return None


class EnrollRequest(BaseModel):
    node_id: str
    name: Optional[str] = None
    lat: float = 0.0
    lon: float = 0.0


class EnrollResponse(BaseModel):
    node_id: str
    pickup_secret: str
    approval_status: str


@router.post("/enroll", response_model=EnrollResponse)
async def enroll_node(
    body: EnrollRequest,
    request: Request,
    x_enrollment_token: Optional[str] = Header(None),
):
    client_ip = request.client.host if request.client else "unknown"
    if not _enroll_limiter.allow(client_ip):
        raise HTTPException(429, "Too many enrollment attempts. Try again later.")

    if not x_enrollment_token:
        logger.warning(f"Enroll 401: missing X-Enrollment-Token from {client_ip} for node_id={body.node_id!r}")
        raise HTTPException(401, "Invalid or missing X-Enrollment-Token")

    org_id = await _resolve_org_for_token(x_enrollment_token)
    if not org_id:
        logger.warning(f"Enroll 401: bad enrollment token from {client_ip} for node_id={body.node_id!r}")
        raise HTTPException(401, "Invalid or missing X-Enrollment-Token")

    node_id = body.node_id.strip()
    if not node_id:
        raise HTTPException(400, "node_id is required")

    existing = await fstore.doc_get("nodes", node_id)

    # -------------------------------------------------------------------
    # CRITICAL GUARD — do not remove or weaken this check.
    #
    # An already-approved node_id must NEVER get a fresh pickup_secret off
    # the fleet-wide enrollment token alone. The fleet token is shared by
    # every node (it ships in every node's .env / setup.sh prompt), so it
    # is the credential most likely to leak. Without this guard, a leaked
    # fleet token plus a guessable node_id (node-001, node-002, ...) would
    # let an attacker "re-enroll" a live, already-approved node and race
    # the real node to GET /nodes/{id}/credentials — stealing its actual
    # api_key before the legitimate device ever asks.
    #
    # Recovery for an approved node goes through the existing admin-only
    # POST /nodes/{id}/reissue-key instead (routers/nodes.py), which
    # requires a Firebase admin token, not the fleet token.
    # -------------------------------------------------------------------
    if existing and existing.get("approval_status") == "approved":
        logger.warning(
            f"Enroll refused: node_id={node_id!r} is already approved — "
            f"refusing to issue a new pickup_secret from the fleet token alone "
            f"(source_ip={client_ip})"
        )
        raise HTTPException(
            403,
            "Node is already approved. This endpoint cannot re-issue credentials "
            "for an approved node from the enrollment token alone — use admin "
            "key reissue.",
        )

    pickup_secret = secrets.token_hex(24)
    doc = {
        "node_id": node_id,
        # A node re-enrolling keeps whatever org_id it already has rather
        # than letting a differently-scoped token silently reassign its
        # tenancy — only a brand-new node_id (or one that predates tenancy
        # entirely) picks up org_id from the token used here.
        "org_id": (existing or {}).get("org_id") or org_id,
        "name": body.name or (existing or {}).get("name") or node_id,
        "lat": body.lat or (existing or {}).get("lat", 0.0),
        "lon": body.lon or (existing or {}).get("lon", 0.0),
        "approval_status": (existing or {}).get("approval_status", "pending"),
        "pickup_secret_hash": _hash_secret(pickup_secret),
    }
    # approval_status stays "pending" for a brand-new node; if it's an
    # existing "pending" or "rejected" node re-enrolling (e.g. lost its
    # pickup_secret before an admin ever approved it), leave whatever
    # status it already has rather than silently flipping "rejected" back
    # to "pending" — that decision belongs to an admin, not this endpoint.
    await fstore.doc_set("nodes", node_id, doc, merge=True)
    logger.info(f"Node enrolled: {node_id} (status={doc['approval_status']}, source_ip={client_ip})")

    return EnrollResponse(node_id=node_id, pickup_secret=pickup_secret, approval_status=doc["approval_status"])


class CredentialsResponse(BaseModel):
    approval_status: str
    api_key: Optional[str] = None


@router.get("/{node_id}/credentials", response_model=CredentialsResponse)
async def get_node_credentials(node_id: str, x_pickup_secret: Optional[str] = Header(None)):
    if not x_pickup_secret:
        raise HTTPException(401, "Missing X-Pickup-Secret header")

    node = await fstore.doc_get("nodes", node_id)
    if not node or not node.get("pickup_secret_hash"):
        raise HTTPException(404, "Unknown node, or node was never enrolled via POST /nodes/enroll")

    if not secrets.compare_digest(_hash_secret(x_pickup_secret), node["pickup_secret_hash"]):
        raise HTTPException(401, "Invalid pickup secret")

    approval_status = node.get("approval_status", "pending")
    if approval_status != "approved":
        return CredentialsResponse(approval_status=approval_status)

    key_doc = await fstore.doc_get("node_keys", node_id)
    if not key_doc or not key_doc.get("api_key"):
        # Approved but no key yet — shouldn't normally happen, approve_node()
        # always writes node_keys in the same call that sets approved. Treat
        # it as "keep polling" rather than erroring the node's retry loop.
        return CredentialsResponse(approval_status=approval_status)

    return CredentialsResponse(approval_status=approval_status, api_key=key_doc["api_key"])