llm-ingestion-pipeline-security/src/llm_ingestion_guard/contract.py

"""contract — the write-time quarantine asserters (BRIEF §5/§6, PLAN §93).

The *differentiator*. Where the other modules detect and report, these functions
enforce an invariant and **raise**: they are the write-time analogue of runtime
least-privilege. They harden the call a pipeline makes around the quarantined,
tool-less transform — the library itself makes no model call.

Three asserts, matching the reusable-contract checklist (BRIEF §6, steps 3-4):

* **(a) tool-less transform** — :func:`assert_tool_less` verifies the model
  request carries zero tools/functions/MCP servers. A successful injection then
  has nothing to act with (checklist step 3).
* **(b) per-stage credential allowlist** — :func:`assert_credential_allowlist`
  verifies the process env holds no credential beyond the ones this stage is
  entitled to: the enrichment stage sees only the model key, never the publish
  credential (checklist step 4).
* **(c) capability isolation** — :func:`scoped_env` produces a correctly scoped
  env so a hijacked stage cannot even read a credential it was never granted;
  the assert then passes by construction.

A fourth asserter lives here for the same reason — it enforces rather than
reports — though it belongs to the transform path rather than the quarantine
checklist: **(d) bounded transform input**, :func:`assert_within_input_cap`,
which the three content-returning surfaces call so an oversize document is
refused instead of half-transformed.

Reference: ``claude-code-llm-wiki`` ``tools/wiki_ingest/enrich.py``
``assert_quarantine`` — a pipeline-specific quarantine gate, generalized here
into framework-agnostic, reusable pieces.

Two safety choices are deliberate and inherited from the reference:

* Credential detection is **name-based**; env *values* are never read into the
  assert. Value-scanning is :mod:`output`'s job for emitted text, not the env.
* A raised :class:`ContractViolation` names only the offending *key* — never a
  credential value — so the exception is safe to route to an alert channel
  (minimal-alert payloads, BRIEF §6, step 8).
"""
from __future__ import annotations

import re
from typing import Iterable, Mapping


class ContractViolation(Exception):
    """A write-time security-contract invariant was violated.

    Carries a machine-readable :attr:`code` and a :attr:`details` tuple of the
    offending identifiers (tool-surface keys or env credential names) so a caller
    can build a minimal, content-free alert payload. The message and details
    never contain a credential value.
    """

    def __init__(self, message: str, *, code: str, details: tuple[str, ...] = ()):
        super().__init__(message)
        self.code = code
        self.details = details


class OversizeInputError(ContractViolation):
    """A transform surface was handed more text than it will transform.

    A *subclass*, not a sibling: a pipeline that already brackets its quarantined
    stage in ``except ContractViolation`` keeps failing closed rather than
    meeting an exception type it has never heard of. :attr:`details` names the
    surface that refused; the sizes go in the message, and neither carries any
    of the input, so the exception stays alert-routable like its parent.
    """


def assert_within_input_cap(text: str, *, surface: str, max_input_chars: int) -> None:
    """Raise :class:`OversizeInputError` unless ``text`` fits the transform cap.

    The write-time asserter for the *transform* surfaces (d). Where the scanners
    bound their work by truncating — findings are lossy in the tail and nothing
    else — a transform returns content, so a prefix-only result is either silent
    data loss or an untransformed tail the payload can be positioned into. The
    invariant these three keep instead is: returned text is always fully
    transformed, or not returned at all.

    ``max_input_chars`` is the largest accepted size, not the smallest rejected
    one.
    """
    size = len(text)
    if size > max_input_chars:
        raise OversizeInputError(
            f"{surface}: input {size} chars exceeds cap {max_input_chars}; "
            "refused rather than partially transformed",
            code="oversize-input",
            details=(surface,),
        )


# --- (a) tool-less transform -----------------------------------------------

# Populated tool surface across Anthropic + OpenAI request shapes. An empty
# list / None is genuinely tool-less; only a *populated* key is a violation.
_TOOL_KEYS = ("tools", "functions", "tool_choice", "function_call", "mcp_servers")


def assert_tool_less(request: Mapping[str, object]) -> None:
    """Assert the model request carries no tool surface (checklist step 3).

    Raises :class:`ContractViolation` (code ``"tool_present"``) if any of
    ``tools``, ``functions``, ``tool_choice``, ``function_call`` or
    ``mcp_servers`` is present *and populated*. ``request`` is the request
    kwargs mapping — framework-agnostic, no SDK type required.
    """
    offending = tuple(sorted(k for k in _TOOL_KEYS if request.get(k)))
    if offending:
        raise ContractViolation(
            "model request carries a tool surface: %s" % ", ".join(offending),
            code="tool_present",
            details=offending,
        )


# --- (b) per-stage credential allowlist ------------------------------------

# Credential-shaped env-var names, matched on the underscore-delimited word
# boundary so ``SECRETARY_NAME`` does not false-positive. Extends the reference
# pipeline's ``CREDENTIAL_NAME_RE`` with a bare ``KEY`` alternative so a
# provider key like ``STRIPE_KEY`` is caught rather than silently missed. This
# is deliberately broad: it also flags non-secret keys (``PARTITION_KEY``,
# ``SORT_KEY``), but that is a LOUD false positive (raises; fixed with a
# one-line allowlist entry), chosen over a silent fail-open on a new provider's
# ``<X>_KEY`` — for an isolation control, fail-loud beats fail-silent (§4.7).
CREDENTIAL_NAME_RE = re.compile(
    r"(^|_)(API_?KEY|APIKEY|KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIALS?)(_|$)"
)


def credential_env_names(env: Mapping[str, str]) -> tuple[str, ...]:
    """Return the sorted env-var names that look like credentials.

    Name-based only: values are never read. ``PATH``/``HOME``/``LANG`` and other
    non-credential vars are ignored.
    """
    return tuple(sorted(
        name for name in env if CREDENTIAL_NAME_RE.search(name.upper())
    ))


def assert_credential_allowlist(env: Mapping[str, str], allowed: Iterable[str]) -> None:
    """Assert the env holds no credential beyond ``allowed`` (checklist step 4).

    Least-privilege: the credential set must be a *subset* of ``allowed`` —
    holding fewer than allowed is safe; holding even one more raises
    :class:`ContractViolation` (code ``"credential_leak"``). The raised message
    names only the offending keys, never their values.
    """
    allowlist = set(allowed)
    leaked = tuple(
        name for name in credential_env_names(env) if name not in allowlist
    )
    if leaked:
        raise ContractViolation(
            "env carries off-allowlist credential(s): %s" % ", ".join(leaked),
            code="credential_leak",
            details=leaked,
        )


# --- (c) capability isolation (env scoping) --------------------------------


def scoped_env(env: Mapping[str, str], allowed: Iterable[str]) -> dict[str, str]:
    """Return a copy of ``env`` with off-allowlist credentials removed.

    Non-credential vars (``PATH`` etc.) pass through so the scoped env still
    works for a subprocess; only credential-shaped keys outside ``allowed`` are
    dropped. The input is not mutated. By construction,
    ``assert_credential_allowlist(scoped_env(env, a), a)`` always passes — a
    hijacked stage cannot read a credential it was never granted.
    """
    allowlist = set(allowed)
    leaked = set(credential_env_names(env)) - allowlist
    return {name: value for name, value in env.items() if name not in leaked}