# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/under-what-conditions-must-an-instance-treat-prompt-injection-taint-as-non-transitive-across-tool-invocation-hops.md
# Remnant Warrant:    https://corpusfetch.com/c/under-what-conditions-must-an-instance-treat-prompt-injection-taint-as-non-transitive-across-tool-invocation-hops.warrant.json
# Markdown SHA-256:   58917e106b7e80d11521a1c10173faa135364bfe03e6996e3ce55722ac93b88a
# Code AST SHA-256:   81725d0432c049f7dbc9b6faf48f234ab37a8152aee233023d776fb3c008e739
# Verified Classes:   TaintPropagationError, HopKind, Hop, Label
# Defensive Guards:   4 fail-closed invariants (zero external deps)

from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Iterable


class TaintPropagationError(RuntimeError):
    """Raised when a hop cannot soundly propagate taint under union semantics."""


class HopKind(Enum):
    PURE = "pure"                 # deterministic function of declared args only
    AMBIENT_READ = "ambient"      # reads mutable state not passed as an argument
    INTERPRETER = "interpreter"   # evaluates its inputs as instructions/code
    LABEL_DROPPING = "dropping"   # persists to a store that loses labels


@dataclass(frozen=True)
class Hop:
    name: str
    kind: HopKind
    declared_ambient: frozenset[str] = field(default_factory=frozenset)


@dataclass(frozen=True)
class Label:
    untrusted: bool
    reasons: frozenset[str] = field(default_factory=frozenset)


def propagate(path: Iterable[Hop], input_label: Label) -> Label:
    """Sound union-propagation, or raise if the path is taint-amplifying."""
    label = input_label
    for hop in path:
        if hop.kind is HopKind.PURE:
            continue
        if hop.kind is HopKind.AMBIENT_READ:
            # Ambient reads are tainted by whatever the ambient state holds.
            # Only sound if the read set is declared and disjoint from mutable
            # attacker-writable state; the runtime cannot verify that here.
            if not hop.declared_ambient:
                raise TaintPropagationError(
                    f"hop {hop.name!r} performs an undeclared ambient read; "
                    "output taint is under-approximated by input-label union"
                )
            label = Label(
                untrusted=True,
                reasons=label.reasons | {f"ambient:{hop.name}"},
            )
            continue
        if hop.kind is HopKind.INTERPRETER:
            # Cannot recover which output token derived from which input token.
            # The only sound policy is pre-hop sanitization, enforced upstream.
            raise TaintPropagationError(
                f"hop {hop.name!r} interprets its inputs; post-hoc label "
                "propagation is unsound, sanitize before the hop"
            )
        if hop.kind is HopKind.LABEL_DROPPING:
            raise TaintPropagationError(
                f"hop {hop.name!r} crosses a label-dropping boundary; "
                "treat every read from it as untrusted"
            )
    return label


def gate(sink: str, label: Label, allow_untrusted: frozenset[str] = frozenset()) -> None:
    if label.untrusted and sink not in allow_untrusted:
        raise TaintPropagationError(
            f"refusing sink {sink!r}: tainted by {sorted(label.reasons)}"
        )