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

import secrets
from typing import Optional
from fastapi import APIRouter, HTTPException, Depends, Query
from pydantic import BaseModel
from app.models import CommandPayload
from app.internal import firestore as fstore
from app.internal.mqtt_handler import mqtt_handler
from app.internal import dynsec
from app.internal.logger import logger
from app.internal.auth import (
    require_admin_token,
    require_service_key_or_admin,
    require_service_or_firebase_token,
    resolve_caller_org_id,
)
from app.routers.tokens import assign_token, release_token

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


@router.get("")
async def list_nodes(decoded: dict = Depends(require_service_or_firebase_token)):
    org_id = await resolve_caller_org_id(decoded)
    if org_id is None:  # service key or platform admin — unrestricted, matches prior behaviour
        return await fstore.collection_list("nodes")
    return await fstore.collection_list("nodes", org_id=org_id)


@router.get("/{node_id}")
async def get_node(node_id: str, decoded: dict = Depends(require_service_or_firebase_token)):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")
    org_id = await resolve_caller_org_id(decoded)
    if org_id is not None and node.get("org_id") != org_id:
        raise HTTPException(404, f"Node '{node_id}' not found.")
    return node


@router.post("/{node_id}/approve")
async def approve_node(node_id: str, _: dict = Depends(require_admin_token)):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")

    api_key = secrets.token_hex(32)

    # dynsec FIRST, Firestore second: if the broker rejects/never confirms
    # the new client, we must not tell Firestore (and the admin UI) the
    # node is approved with a key mosquitto doesn't actually recognise —
    # that's exactly the silent-drift the two-sources-of-truth problem
    # warns about. See app/internal/dynsec.py.
    try:
        await dynsec.upsert_node_client(node_id, api_key)
    except dynsec.DynsecError as e:
        logger.error(f"Approve {node_id!r}: dynsec upsert failed, NOT writing Firestore: {e}")
        raise HTTPException(502, f"Could not provision MQTT credentials for node: {e}")

    await fstore.doc_set("node_keys", node_id, {"node_id": node_id, "api_key": api_key}, merge=False)
    await fstore.doc_update("nodes", node_id, {"approval_status": "approved"})
    # TODO(mqtt-cutover): drop this MQTT push once nodes pull their key via
    # GET /nodes/{id}/credentials (routers/enrollment.py) exclusively — see
    # MQTT-PUBLIC-AUTH-PLAN.md "Rollout order" step 6. Kept for node-26.
    mqtt_handler.publish_node_key(node_id, api_key)
    return {"ok": True}


@router.delete("/{node_id}", status_code=204)
async def delete_node(node_id: str, _: dict = Depends(require_admin_token)):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")
    try:
        await dynsec.delete_node_client(node_id)
    except dynsec.DynsecError as e:
        logger.error(f"Delete {node_id!r}: dynsec deleteClient failed, NOT deleting Firestore docs: {e}")
        raise HTTPException(502, f"Could not revoke MQTT credentials for node: {e}")
    await fstore.doc_delete("node_keys", node_id)
    await fstore.doc_delete("nodes", node_id)


@router.post("/{node_id}/reject")
async def reject_node(node_id: str, _: dict = Depends(require_admin_token)):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")
    await fstore.doc_update("nodes", node_id, {"approval_status": "rejected"})
    return {"ok": True}


@router.post("/{node_id}/command")
async def send_command(
    node_id: str,
    cmd: CommandPayload,
    _: dict = Depends(require_service_key_or_admin),
):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")

    payload = cmd.model_dump(exclude_none=True)

    if cmd.action == "discord_join":
        # Resolve system doc once — used for preferred token and presence name.
        system_doc = None
        system_id = node.get("assigned_system_id")
        if system_id:
            system_doc = await fstore.doc_get_cached("systems", system_id)

        # Explicit preferred_token_id in the request beats the system-level preference.
        preferred = payload.pop("preferred_token_id", None) or (system_doc or {}).get("preferred_token_id")
        token = await assign_token(node_id, preferred_token_id=preferred)
        if not token:
            raise HTTPException(503, "No Discord bot tokens available in the pool.")
        payload["token"] = token

        # Pass system name so the bot can set its Discord presence on join.
        system_name = (system_doc or {}).get("name")
        if system_name:
            payload["system_name"] = system_name

    elif cmd.action == "discord_leave":
        await release_token(node_id)

    if not mqtt_handler.send_command(node_id, payload):
        raise HTTPException(503, "MQTT broker unavailable — command not delivered.")
    return {"ok": True}


@router.post("/{node_id}/reissue-key")
async def reissue_node_key(node_id: str, _: dict = Depends(require_admin_token)):
    """Generate a new API key for the node and push it via MQTT (retained).
    Use this to rotate a key or recover a node whose key was lost."""
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")
    api_key = secrets.token_hex(32)
    try:
        await dynsec.upsert_node_client(node_id, api_key)
    except dynsec.DynsecError as e:
        logger.error(f"Reissue {node_id!r}: dynsec upsert failed, NOT writing Firestore: {e}")
        raise HTTPException(502, f"Could not update MQTT credentials for node: {e}")
    await fstore.doc_set("node_keys", node_id, {"node_id": node_id, "api_key": api_key}, merge=False)
    # TODO(mqtt-cutover): drop this MQTT push once nodes pull their key via
    # GET /nodes/{id}/credentials (routers/enrollment.py) exclusively — see
    # MQTT-PUBLIC-AUTH-PLAN.md "Rollout order" step 6. Kept for node-26.
    mqtt_handler.publish_node_key(node_id, api_key)
    return {"ok": True}


@router.post("/{node_id}/config/{system_id}")
async def assign_system(
    node_id: str,
    system_id: str,
    hardware_preset: str = Query("rtl-sdr-v3"),
    ppm_override: Optional[float] = Query(None),
    _: dict = Depends(require_service_key_or_admin),
):
    """
    Assign a system to a node. Fetches the system config from Firestore
    and pushes it to the node via MQTT, then marks the node as configured.
    """
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")

    system = await fstore.doc_get("systems", system_id)
    if not system:
        raise HTTPException(404, f"System '{system_id}' not found.")

    # Include hardware preset, node type, and enforce timeout in the push
    push_payload = {
        **system, 
        "hardware_preset": hardware_preset,
        "node_type": node.get("node_type", "fixed"),
        "enforce_override_timeout": node.get("enforce_override_timeout", True),
    }
    if ppm_override is not None:
        push_payload["ppm_override"] = ppm_override
    mqtt_handler.push_config(node_id, push_payload)

    # Update Firestore
    node_updates = {
        "assigned_system_id": system_id,
        "configured": True,
        "hardware_preset": hardware_preset,
    }
    if ppm_override is not None:
        node_updates["ppm_override"] = ppm_override
    await fstore.doc_update("nodes", node_id, node_updates)

    return {"ok": True}


class NodeUpdateBody(BaseModel):
    node_type: Optional[str] = None
    enforce_override_timeout: Optional[bool] = None


@router.patch("/{node_id}")
async def update_node(
    node_id: str,
    body: NodeUpdateBody,
    _: dict = Depends(require_admin_token),
):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")

    updates = body.model_dump(exclude_unset=True)
    if not updates:
        return {"ok": True}

    await fstore.doc_update("nodes", node_id, updates)

    # Re-push config to apply new node settings locally
    updated_node = await fstore.doc_get("nodes", node_id)
    assigned_system_id = updated_node.get("assigned_system_id")
    if assigned_system_id:
        system = await fstore.doc_get("systems", assigned_system_id)
        if system:
            push_payload = {
                **system,
                "hardware_preset": updated_node.get("hardware_preset", "rtl-sdr-v3"),
                "node_type": updated_node.get("node_type", "fixed"),
                "enforce_override_timeout": updated_node.get("enforce_override_timeout", True),
            }
            if updated_node.get("ppm_override") is not None:
                push_payload["ppm_override"] = updated_node["ppm_override"]
            mqtt_handler.push_config(node_id, push_payload)

    return {"ok": True}


class AckOverrideBody(BaseModel):
    timeout_minutes: int = 1440


@router.post("/{node_id}/override/ack")
async def ack_override(
    node_id: str,
    body: AckOverrideBody,
    _: dict = Depends(require_service_key_or_admin),
):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")

    from datetime import datetime, timezone, timedelta
    new_timeout = datetime.now(timezone.utc) + timedelta(minutes=body.timeout_minutes)
    
    await fstore.doc_update("nodes", node_id, {
        "override_timeout_at": new_timeout.isoformat()
    })
    return {"ok": True, "override_timeout_at": new_timeout.isoformat()}


@router.post("/{node_id}/override/reset")
async def reset_override(
    node_id: str,
    _: dict = Depends(require_service_key_or_admin),
):
    node = await fstore.doc_get("nodes", node_id)
    if not node:
        raise HTTPException(404, f"Node '{node_id}' not found.")

    assigned_system_id = node.get("assigned_system_id")
    if assigned_system_id:
        system = await fstore.doc_get("systems", assigned_system_id)
        if system:
            mqtt_handler.push_config(node_id, system)

    await fstore.doc_update("nodes", node_id, {
        "is_overridden": False,
        "override_system_id": None,
        "override_timeout_at": None,
    })
    return {"ok": True}