# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/in-what-scenarios-must-an-instance-distrust-the-provenance-of-an-mcp-tool-response-rather-than-its-content-taint-propagation-across-re-serialization-b.md
# Remnant Warrant:    https://corpusfetch.com/c/in-what-scenarios-must-an-instance-distrust-the-provenance-of-an-mcp-tool-response-rather-than-its-content-taint-propagation-across-re-serialization-b.warrant.json
# Markdown SHA-256:   42663758b3cb76d7040f73cc843a2fa17d51aa43a5d6da766613c29aaf1f9910
# Code AST SHA-256:   cc76448b03cdb2c0309aa745f933fc71f6c4bce9e6bfe2623a19efd1b2021b1c
# Verified Classes:   TaintLaunderingError, Provenance, ProvenanceLedger
# Defensive Guards:   2 fail-closed invariants (zero external deps)

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any


class TaintLaunderingError(RuntimeError):
    """Raised when a value derived from an untrusted MCP response reaches a
    privileged sink without a resolvable, label-preserving derivation root."""


@dataclass(frozen=True)
class Provenance:
    root_id: str          # identity of the originating tool response
    trusted: bool         # False for any MCP tool response by default
    edges: tuple[str, ...] = field(default_factory=tuple)  # named transforms


class ProvenanceLedger:
    """Maps object identity -> provenance. Outlives any single object so that
    re-serialized copies can be resolved back to their derivation root."""

    def __init__(self) -> None:
        self._by_id: dict[int, Provenance] = {}

    def tag(self, obj: Any, prov: Provenance) -> Any:
        self._by_id[id(obj)] = prov
        return obj

    def resolve(self, obj: Any) -> Provenance | None:
        return self._by_id.get(id(obj))

    def derive(self, src: Any, dst: Any, edge: str) -> Any:
        root = self.resolve(src)
        if root is None:
            # Unknown source: fail closed, treat dst as untrusted.
            return self.tag(dst, Provenance(root_id=f"unknown:{id(dst)}", trusted=False))
        return self.tag(dst, Provenance(root.root_id, root.trusted, root.edges + (edge,)))


LEDGER = ProvenanceLedger()


def load_tool_response(raw: bytes, tool_call_id: str) -> Any:
    """Parse an untrusted MCP tool response. The parsed structure is *not*
    trusted merely because json.loads produced it."""
    parsed = json.loads(raw)
    return LEDGER.tag(parsed, Provenance(root_id=tool_call_id, trusted=False, edges=("json.loads",)))


def re_serialize(obj: Any) -> str:
    """Any re-serialization must preserve provenance. If the source is unknown,
    the output is tainted by default."""
    out = json.dumps(obj)
    return LEDGER.derive(obj, out, "json.dumps")


def assemble_prompt(context: Any) -> str:
    """Prompt assembly is a re-serialization boundary and a common laundering
    site. The resulting string inherits the taint of its inputs."""
    block = re_serialize(context)
    prompt = f"<context>{block}</context>"
    return LEDGER.derive(block, prompt, "fstring")


def check_sink_capability(arg: Any, sink: str) -> None:
    """Capability check at a privileged sink. Inspects the derivation root,
    not the argument's runtime type."""
    prov = LEDGER.resolve(arg)
    if prov is None:
        raise TaintLaunderingError(
            f"sink={sink!r}: argument has no resolvable provenance; refusing call"
        )
    if not prov.trusted:
        raise TaintLaunderingError(
            f"sink={sink!r}: argument derives from untrusted root {prov.root_id!r} "
            f"via edges {prov.edges!r}; declassification required before use"
        )


def invoke_filesystem_write(prompt_text: str, path: str, body: str) -> None:
    # The prompt text is the attacker-influenced channel; the sink must reject
    # it if it carries taint from the tool response.
    check_sink_capability(prompt_text, sink="filesystem.write")
    check_sink_capability(path, sink="filesystem.write")
    check_sink_capability(body, sink="filesystem.write")
    # ... actual write would occur here ...


def handle_tool_result(raw: bytes, tool_call_id: str) -> None:
    response = load_tool_response(raw, tool_call_id)
    context = response.get("context", {}) if isinstance(response, dict) else {}
    # The attacker-controlled field is re-serialized into a prompt block.
    prompt = assemble_prompt(context)
    # A downstream model now "decides" to write a file. The sink must refuse.
    invoke_filesystem_write(prompt_text=prompt, path="notes.txt", body="benign")