Pi lifecycle SAS-after-enrollment - temporary deploy

enrollment_wire.py

"""Wiring glue between MQTT _on_message and the enrollment_handler/apply pair.

This module is the ONE entry point invoked by ``mqtt/client/thread.py``
when an ``discovery/{device_uid}/enrollment`` message arrives. It:

1. Decodes the JSON payload (rejects non-object).
2. Loads the Pi's persistent state (or builds an ephemeral fallback).
3. Auto-transitions BLANK → BOOTSTRAPPED → DISCOVERABLE if needed, so a
   freshly imaged Pi receiving its first enrollment can accept it
   (state.json is wiped by smc-clone-hygiene on every boot until we
   complete the factory image cleanup sprint).
4. Calls ``handle_discovery_enrollment`` to validate + plan.
5. On valid plan: invokes ``apply_enrollment_plan`` with concrete
   publish/save_state callbacks bound to the MQTT client + state_store.
6. Returns a ``DispatchResult`` so thread.py only does a print + state.push.

All exceptions are caught and reported in the result — thread.py never
sees an uncaught exception from enrollment handling.
"""
from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Callable

logger = logging.getLogger(__name__)


@dataclass
class DispatchResult:
    """Outcome of dispatch_enrollment_topic() — for logging by thread.py."""
    handled: bool  # True if topic matched + dispatch attempted
    valid_plan: bool = False
    apply_success: bool = False
    error_step: str = ""
    error_message: str = ""
    node_id: str = ""
    site_id: str = ""
    sas_lp_73_result: str = "not_attempted"


def dispatch_enrollment_topic(
    topic: str,  # noqa: ARG001 — kept for future logging / metrics by topic
    payload: str,
    *,
    publish_mqtt: Callable[[str, str], Any],
    device_uid_resolver: Callable[[], str],
    bootstrap_secret_resolver: Callable[[], bytes],
) -> DispatchResult:
    """Dispatch a single ``discovery/+/enrollment`` MQTT message.

    Args:
        topic: full MQTT topic (must end with ``/enrollment``).
        payload: raw JSON payload string.
        publish_mqtt: ``(topic, payload_json_str) -> None``. Used to publish
            the receipt ack — caller controls QoS/retain semantics.
        device_uid_resolver: ``() -> str``. Lazy read of DEVICE_UID at
            dispatch time (so it picks up updates from runtime.env).
        bootstrap_secret_resolver: ``() -> bytes``. Same pattern.

    Returns DispatchResult. ``handled=False`` only when the JSON is malformed.
    """
    # Lazy imports — keep this module light when not invoked
    from backend.lifecycle.enrollment_apply import apply_enrollment_plan
    from backend.lifecycle.enrollment_handler import handle_discovery_enrollment
    from backend import state_store

    result = DispatchResult(handled=False)

    # 1) JSON parse
    try:
        if not payload or not payload.strip().startswith("{"):
            result.error_step = "payload_parse"
            result.error_message = "payload is not a JSON object"
            return result
        env_obj = json.loads(payload)
    except Exception as e:
        result.error_step = "payload_parse"
        result.error_message = f"json decode: {e}"
        return result
    if not isinstance(env_obj, dict):
        result.error_step = "payload_parse"
        result.error_message = "payload is not a JSON object"
        return result

    result.handled = True

    # 2) Resolve identity + bootstrap secret
    try:
        device_uid = device_uid_resolver() or ""
        bootstrap_secret = bootstrap_secret_resolver()
    except Exception as e:
        result.error_step = "resolve_identity"
        result.error_message = str(e)
        return result

    # 3) Load state.json (fallback ephemeral on missing/corrupt file).
    # Pass DEFAULT_PATH explicitly so tests monkeypatching the module attr
    # see their override (Python evaluates `path=DEFAULT_PATH` at def-time,
    # not call-time).
    try:
        state_data = state_store.init_state(path=state_store.DEFAULT_PATH)
    except Exception:
        state_data = state_store.StateData()

    # 4) Auto-bootstrap to DISCOVERABLE if Pi is fresh (BLANK).
    # Cheap workaround: smc-clone-hygiene wipes state.json on every boot
    # until the factory-image-cleanup sprint lands. Without this step the
    # first enrollment after a wipe always fails on "cannot transition
    # blank -> enrolling".
    if state_data.state == state_store.PiState.BLANK:
        try:
            state_store.transition_to(
                state_data, state_store.PiState.BOOTSTRAPPED,
                "auto:on_enrollment_received",
            )
            state_store.transition_to(
                state_data, state_store.PiState.DISCOVERABLE,
                "auto:on_enrollment_received",
            )
        except Exception:
            pass  # if the auto-bootstrap fails, handle will reject below

    # 5) Validate + plan
    plan = handle_discovery_enrollment(
        env_obj, bootstrap_secret, state_data, device_uid=device_uid,
    )
    if not plan.valid:
        result.error_step = "plan_invalid"
        result.error_message = f"{plan.error_code}: {plan.error_message}"
        return result

    result.valid_plan = True
    result.node_id = plan.node_id
    result.site_id = plan.site_id

    # 6) Build apply callbacks bound to MQTT + state_store
    def _publish_ack(envelope: dict) -> None:
        ack_topic = f"cms/ack/{device_uid}"
        publish_mqtt(ack_topic, json.dumps(envelope, separators=(",", ":")))

    def _save_state(sd: Any) -> None:
        # Explicit path same reason as init_state: Python evaluates the
        # default at def-time, so tests monkeypatching the module attr
        # see their override only when we re-resolve at call-time.
        state_store.save_state(sd, path=state_store.DEFAULT_PATH)

    # 7) Apply side effects
    apply_result = apply_enrollment_plan(
        plan, state_data, _publish_ack, _save_state,
    )
    result.apply_success = apply_result.success
    result.sas_lp_73_result = apply_result.sas_lp_73_result
    if not apply_result.success:
        result.error_step = apply_result.error_step
        result.error_message = apply_result.error_message
        return result

    # 8) Trigger SAS init cycle in background (AC #2.3).
    # Fire-and-forget on the firmware's main asyncio loop. The coroutine
    # awaits sas_bus_submit("run_init_sequence"), transitions state to
    # ENROLLED on success, and publishes a second ack (signed with the new
    # secret_pi) so the CMS can mark the device ACTIVATED. Errors are caught
    # so a SAS bus failure doesn't propagate to the MQTT dispatcher.
    _trigger_sas_init_after_enrollment(plan, device_uid)

    return result


def _trigger_sas_init_after_enrollment(plan, device_uid: str) -> None:
    """Schedule the SAS LP 0x73 cycle + ENROLLED transition + ack publish.

    Runs in the firmware's main asyncio loop via run_coroutine_threadsafe
    (the caller is the MQTT paho thread). Best-effort: any error is logged
    and the device stays in ENROLLING until the next admin intervention.
    """
    from backend.lifecycle.envelope import build_signed_envelope
    from backend import state_store

    secret_pi_new = (plan.secret_pi or "").encode()
    cmd_id = plan.cmd_id or ""
    node_id = plan.node_id or ""

    try:
        from backend import state as _backend_state
        loop = getattr(_backend_state, "MAIN_LOOP", None)
        if loop is None or not loop.is_running():
            logger.warning(
                "main asyncio loop not available — SAS init skipped "
                "(device stays in ENROLLING until manual sas_init trigger)"
            )
            return

        async def _sas_init_coro():
            ok = False
            err_str = ""
            try:
                from backend import app as _app
                rid = f"enrollment-{cmd_id[:8]}"
                logger.info("triggering sas_bus_submit run_init_sequence rid=%s", rid)
                res = await _app.sas_bus_submit(
                    "run_init_sequence",
                    priority=5,
                    payload={"rid": rid},
                )
                ok = bool((res or {}).get("ok"))
                if not ok:
                    err_str = str((res or {}).get("err") or (res or {}).get("error") or "sas_init_failed")
            except Exception as e:
                logger.error("sas_bus_submit run_init_sequence failed: %s", e, exc_info=True)
                err_str = f"{type(e).__name__}: {e}"

            # On success: transition state to ENROLLED + save
            if ok:
                try:
                    sd = state_store.init_state(path=state_store.DEFAULT_PATH)
                    state_store.transition_to(
                        sd, state_store.PiState.ENROLLED,
                        "sas.aft_register_completed",
                    )
                    state_store.save_state(sd, path=state_store.DEFAULT_PATH)
                    logger.info("state transitioned to ENROLLED for node_id=%s", node_id)
                except Exception as e:
                    logger.error("ENROLLED transition failed: %s", e, exc_info=True)

            # Build + publish 2nd ack signed with the new secret_pi
            try:
                from backend.mqtt.client.publish import mqtt_publish_json
                now_iso = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
                ack_type = "sas.registration.ack" if ok else "enrollment.failed"
                ack_data = {
                    "cmd_id": cmd_id,
                    "result": "applied" if ok else "sas_init_failed",
                    "applied_at": now_iso,
                }
                if not ok:
                    ack_data["error"] = err_str
                ack_env = build_signed_envelope(
                    secret=secret_pi_new,
                    from_=f"device:{node_id}",
                    to_="cms",
                    type_=ack_type,
                    data=ack_data,
                )
                mqtt_publish_json(
                    f"cms/ack/{device_uid}", ack_env, qos=1, retain=False,
                )
                logger.info(
                    "published %s ack ok=%s node_id=%s",
                    ack_type, ok, node_id,
                )
            except Exception as e:
                logger.error("publish 2nd ack failed: %s", e, exc_info=True)

        asyncio.run_coroutine_threadsafe(_sas_init_coro(), loop)
    except Exception as e:
        logger.warning("could not schedule SAS init coroutine: %s", e)


__all__ = [
    "DispatchResult",
    "dispatch_enrollment_topic",
]