stack/packages/mosaic/framework/tools/lease-broker/normative_fragments.py

#!/usr/bin/env python3
"""Fail-closed construction of verbatim-hashed normative fragments.

This is the single construction path for the Claude and Pi adapters.  The
payload contains only versioned source metadata and exact validated source
bytes.  ``h_payload`` is derived afterwards and is deliberately not representable
as a payload input, preventing cryptographic self-reference.
"""

from __future__ import annotations

import hashlib
import hmac
import struct
from collections.abc import Sequence
from typing import Final

MAX_FRAGMENT_BYTES: Final = 64 * 1024
HASH_DOMAIN_SEPARATOR: Final = b"MOSAIC/H_PAYLOAD/v1\x00"
SOURCE_DOMAIN_SEPARATOR: Final = b"MOSAIC/H_SOURCE/v1\x00"
_PAYLOAD_VERSION_LABEL: Final = b"MOSAIC/B_PAYLOAD/v1"


class NormativeFragment:
    """A source identity, its expected digest, and its exact resolved bytes."""

    def __init__(self, source_id: str, content: bytes | None, expected_sha256: str) -> None:
        self.source_id = source_id
        self.content = content
        self.expected_sha256 = expected_sha256


class ConstructionResult:
    """A source-admission decision and, only when admitted, derived payload values.

    ``promotion`` means the construction has produced the only values a later
    receipt protocol may use to attempt promotion. It never performs broker
    promotion itself. A REFUSED result has no payload/hash values, so it cannot
    advance to that later protocol.
    """

    def __init__(
        self,
        injection_decision: str,
        promotion: bool,
        source_reason: str | None,
        b_payload: bytes | None,
        h_payload: str | None,
        h_source: str | None,
    ) -> None:
        self.injectionDecision = injection_decision
        self.promotion = promotion
        self.source_reason = source_reason
        self.b_payload = b_payload
        self.h_payload = h_payload
        self.h_source = h_source


def length_frame(parts: Sequence[bytes]) -> bytes:
    """Encode a finite ordered byte sequence with unambiguous 64-bit framing."""

    framed = bytearray(struct.pack(">Q", len(parts)))
    for part in parts:
        if not isinstance(part, bytes):
            raise TypeError("length framing requires bytes")
        framed.extend(struct.pack(">Q", len(part)))
        framed.extend(part)
    return bytes(framed)


def _sha256_hex(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def _valid_digest(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) == 64
        and all(character in "0123456789abcdef" for character in value)
    )


def _source_reason(fragment: NormativeFragment) -> str | None:
    if not isinstance(fragment.source_id, str) or not fragment.source_id:
        return "missing"
    if fragment.content is None:
        return "missing"
    if not isinstance(fragment.content, bytes):
        return "missing"
    if len(fragment.content) > MAX_FRAGMENT_BYTES:
        return "oversize"
    if not _valid_digest(fragment.expected_sha256):
        return "hash-mismatch"
    actual = _sha256_hex(fragment.content)
    if not hmac.compare_digest(actual, fragment.expected_sha256):
        return "hash-mismatch"
    return None


def _refused(reason: str) -> ConstructionResult:
    return ConstructionResult(
        injection_decision="REFUSED",
        promotion=False,
        source_reason=reason,
        b_payload=None,
        h_payload=None,
        h_source=None,
    )


def build_payload(
    *,
    manifest_version: int,
    generator_version: str,
    fragments: Sequence[NormativeFragment],
) -> ConstructionResult:
    """Construct B_payload then H_payload after fail-closed source validation.

    The sequence order is caller-supplied resolved-source order and is encoded
    directly. Changing source order, source identity, metadata, or any exact
    fragment byte therefore changes the framed B_payload and its derived hash.
    """

    if type(manifest_version) is not int or manifest_version < 0:
        return _refused("missing")
    if not isinstance(generator_version, str) or not generator_version:
        return _refused("missing")

    validated = list(fragments)
    if not validated:
        return _refused("missing")
    for fragment in validated:
        if not isinstance(fragment, NormativeFragment):
            return _refused("missing")
        reason = _source_reason(fragment)
        if reason is not None:
            return _refused(reason)

    source_identities = [
        length_frame([
            fragment.source_id.encode("utf-8"),
            fragment.expected_sha256.encode("ascii"),
        ])
        for fragment in validated
    ]
    h_source = _sha256_hex(SOURCE_DOMAIN_SEPARATOR + length_frame(source_identities))
    payload_parts = [
        _PAYLOAD_VERSION_LABEL,
        str(manifest_version).encode("ascii"),
        generator_version.encode("utf-8"),
        *source_identities,
        *(fragment.content for fragment in validated),
    ]
    # h_payload is intentionally not an argument or field in payload_parts.
    b_payload = length_frame(payload_parts)
    h_payload = _sha256_hex(HASH_DOMAIN_SEPARATOR + length_frame([b_payload]))
    return ConstructionResult(
        injection_decision="ACCEPTED",
        promotion=True,
        source_reason=None,
        b_payload=b_payload,
        h_payload=h_payload,
        h_source=h_source,
    )


def build_for_claude(**kwargs: object) -> ConstructionResult:
    """Claude adapter entrypoint; delegates to the sole shared constructor."""

    return build_payload(**kwargs)  # type: ignore[arg-type]


def build_for_pi(**kwargs: object) -> ConstructionResult:
    """Pi adapter entrypoint; delegates to the sole shared constructor."""

    return build_payload(**kwargs)  # type: ignore[arg-type]