Pi lifecycle defer MQTT creds + enrollment_wire

enrollment_handler.py

"""Handler for discovery.enrollment messages from CMS.

Validates the incoming signed envelope (HMAC + ts window + nonce replay),
extracts the PiEnrollmentPayload fields, and plans the side effects. The
module is pure: it never writes to disk, runs SAS commands, or publishes
MQTT itself. The caller (wiring layer in mqtt/client/thread.py) reads
``EnrollmentPlan`` and performs the apply steps in order.

Apply order expected from the caller (when plan.valid is True):
1. Write runtime.env with plan.runtime_env_updates (atomic, persistent).
2. Publish plan.ack_envelope_pending (result="received") to acknowledge
   receipt before the SAS cycle.
3. Run SAS LP 0x73 cycle with (asset_number, registration_key_hex, pos_id_hex).
4. On SAS success: transition state.json → ENROLLED and publish a
   second ack (result="applied") signed with the NEW secret_pi.
5. On SAS failure: stay in ENROLLING and publish enrollment.failed.
6. state_store.save_state(state_data) at the end (transition already
   applied to state_data by this handler).
"""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Optional

from .envelope import build_signed_envelope, parse_envelope
from .message_validator import validate_incoming_message, ValidationResult
from ..state_store import PiState, record_cmd_ack, transition_to


class EnrollmentHandlerError(Exception):
    pass


@dataclass
class EnrollmentPlan:
    """Result of validating + planning a discovery.enrollment envelope."""
    valid: bool
    error_code: str = ""
    error_message: str = ""

    # Identity to assign at apply time
    node_id: str = ""
    site_id: str = ""
    company_id: int = 0
    location_id: int = 0
    machine_id: int = 0
    aft_key_id_ref: int = 0

    # MQTT credentials (Pi reconnects with these after apply)
    mqtt_user: str = ""
    mqtt_pass: str = ""

    # Per-device secret used for all subsequent signed messages
    secret_pi: str = ""
    applied_version: int = 1

    # SAS LP 0x73 parameters
    sas_asset_number: int = 0
    sas_registration_key_hex: str = ""  # 40 hex chars (20 bytes)
    sas_pos_id_hex: str = ""  # 8 hex chars (4 bytes)

    # Atomic update set for runtime.env
    runtime_env_updates: dict = field(default_factory=dict)

    # State machine transition already applied to state_data
    new_state: Optional[str] = None

    # cmd_id (= envelope.msg_id) used for ack correlation
    cmd_id: str = ""

    # Ack envelope (signed with bootstrap, result="received") — caller must
    # publish this immediately after writing runtime.env. A second ack with
    # result="applied" is built later by the caller once the SAS cycle
    # completes (signed with the new secret_pi).
    ack_envelope_pending: Optional[dict] = None


_MIN_SECRET_PI_LEN = 32
_EXPECTED_TYPE = "discovery.enrollment"
_ACK_TYPE = "discovery.enrollment.ack"


def _build_pending_ack(
    bootstrap_secret: bytes,
    device_uid: str,
    msg_id: str,
) -> dict:
    """Ack 'received' signed with bootstrap secret (Pi has not yet adopted secret_pi)."""
    return build_signed_envelope(
        secret=bootstrap_secret,
        from_=f"device:{device_uid}",
        to_="cms",
        type_=_ACK_TYPE,
        data={
            "msg_id": msg_id,
            "result": "received",
            "received_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        },
    )


def handle_discovery_enrollment(
    envelope: dict,
    bootstrap_secret: bytes,
    state_data: Any,
    *,
    device_uid: str,
) -> EnrollmentPlan:
    """Validate + plan a discovery.enrollment message.

    On valid=False, plan.runtime_env_updates stays empty and the state machine
    is NOT mutated. On valid=True, state_data has already been transitioned to
    ENROLLING — caller must save_state after applying the side effects.
    """
    plan = EnrollmentPlan(valid=False)

    # 1) Envelope-level validation (HMAC + ts window + nonce replay)
    result = validate_incoming_message(envelope, bootstrap_secret, state_data)
    if result != ValidationResult.OK:
        plan.error_code = result.value.upper()
        plan.error_message = f"Validation failed: {result.value}"
        return plan

    # 2) Type check
    parsed = parse_envelope(envelope)
    if parsed.type != _EXPECTED_TYPE:
        plan.error_code = "WRONG_TYPE"
        plan.error_message = f"Expected {_EXPECTED_TYPE}, got {parsed.type}"
        return plan

    # 3) Payload schema (manual — pydantic is too heavy for the Pi runtime)
    data = parsed.data or {}
    sas_reg = data.get("sas_aft_registration")
    if not isinstance(sas_reg, dict):
        plan.error_code = "INVALID_SAS_REG"
        plan.error_message = "sas_aft_registration must be an object"
        return plan

    node_id = str(data.get("node_id", "")).strip()
    secret_pi = str(data.get("secret_pi", ""))
    if not node_id:
        plan.error_code = "MISSING_NODE_ID"
        plan.error_message = "node_id is required"
        return plan
    if len(secret_pi) < _MIN_SECRET_PI_LEN:
        plan.error_code = "INVALID_SECRET_PI"
        plan.error_message = (
            f"secret_pi must be at least {_MIN_SECRET_PI_LEN} chars"
        )
        return plan

    plan.node_id = node_id
    plan.site_id = str(data.get("site_id", ""))
    plan.company_id = int(data.get("company_id", 0) or 0)
    plan.location_id = int(data.get("location_id", 0) or 0)
    plan.machine_id = int(data.get("machine_id", 0) or 0)
    plan.aft_key_id_ref = int(data.get("aft_key_id_ref", 0) or 0)
    plan.mqtt_user = str(data.get("mqtt_user", ""))
    plan.mqtt_pass = str(data.get("mqtt_pass", ""))
    plan.secret_pi = secret_pi
    plan.applied_version = int(data.get("applied_version", 1) or 1)

    plan.sas_asset_number = int(sas_reg.get("asset_number", 0) or 0)
    plan.sas_registration_key_hex = str(sas_reg.get("registration_key", ""))
    plan.sas_pos_id_hex = str(sas_reg.get("pos_id", ""))

    # 4) Build the runtime.env update set.
    # NOTE: MQTT_USER/MQTT_PASS are intentionally NOT written here. The new
    # per-device MQTT credentials issued by the CMS only work if the broker
    # has the user provisioned via DynSec, but the backend's
    # dynsec_admin_publish is currently a stub (mocked publish on
    # $CONTROL/dynamic-security/v1 without an actual DynSec admin client).
    # Until that stub is wired (polish_final), overwriting MQTT_USER here
    # would lock the Pi out of the broker (rc=135 NOT_AUTHORIZED) at the
    # next restart. The Pi keeps its bootstrap MQTT credentials and still
    # signs all subsequent messages with the new secret_pi.
    plan.runtime_env_updates = {
        "NODE_ID": plan.node_id,
        "SITE_ID": plan.site_id,
        "COMPANY_ID": str(plan.company_id),
        "LOCATION_ID": str(plan.location_id),
        "MACHINE_ID": str(plan.machine_id),
        # "MQTT_USER": plan.mqtt_user,  # deferred to dynsec_admin_publish fix
        # "MQTT_PASS": plan.mqtt_pass,  # deferred to dynsec_admin_publish fix
        "SECRET_PI": plan.secret_pi,
        "APPLIED_VERSION": str(plan.applied_version),
        "AFT_KEY_ID_REF": str(plan.aft_key_id_ref),
        "AFT_REGISTRATION_KEY": plan.sas_registration_key_hex,
        "AFT_POS_ID": plan.sas_pos_id_hex,
        "STATE": PiState.ENROLLING.value,
    }

    # 5) Transition state machine to ENROLLING
    try:
        transition_to(state_data, PiState.ENROLLING, "discovery.enrollment")
    except Exception as e:
        plan.error_code = "TRANSITION_FAILED"
        plan.error_message = str(e)
        plan.runtime_env_updates = {}
        return plan
    plan.new_state = PiState.ENROLLING.value

    # 6) Build the receipt ack (caller publishes after writing runtime.env)
    now_iso = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
    plan.cmd_id = parsed.msg_id
    ack = _build_pending_ack(bootstrap_secret, device_uid, parsed.msg_id)
    record_cmd_ack(
        state_data, parsed.msg_id,
        received_at=parsed.ts, ack_sent_at=now_iso,
        ack_msg_id=ack["msg_id"],
    )
    plan.ack_envelope_pending = ack

    plan.valid = True
    return plan


__all__ = [
    "EnrollmentPlan",
    "EnrollmentHandlerError",
    "handle_discovery_enrollment",
]

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",
]