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

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()


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 settings.enrollment_token:
        raise HTTPException(503, "Enrollment is not configured on this server.")
    if not x_enrollment_token or not secrets.compare_digest(x_enrollment_token, settings.enrollment_token):
        logger.warning(f"Enroll 401: bad/missing 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,
        "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"])