llm-ingestion-pipeline-security/tests/inbox_frontend.py

"""OKF inbox front-end — stage 1 of the two-stage upload showcase (PLAN §247).

Reads the files a human actually drops into an inbox and *materializes* them into
an OKF bundle ``{concept_path: document_text}`` with provenance, so the stage-2
guard (:func:`llm_ingestion_guard.okf.import_bundle`) can validate every concept.
This stage owns the container/format threats; the guard owns the text/structural
threats.

**Placement.** This lives in the test tree, not ``src/``: the extraction parsers
are showcase/dev-scoped (``python-docx``/``python-pptx`` in the ``dev`` extra,
never core ``dependencies``), and the shippable core stays stdlib-only. It is an
in-repo demonstration a consumer reads and adapts, not v1 shipped code.

Slice 2a covers the text formats — ``.txt`` and ``.md`` (stdlib only). ``.zip``
(zip-slip / zip-bomb), ``.csv`` (formula injection), folders, ``.docx`` and
``.pptx`` land in later slices.
"""
from __future__ import annotations

import stat
import zipfile
from dataclasses import dataclass
from pathlib import Path

from llm_ingestion_guard.okf import import_bundle, Origin, Channel, BundleResult
from llm_ingestion_guard.disposition import Disposition


# Aggregate disposition -> inbox verdict. Fail-secure by default: a front-end
# refusal (a container threat the guard never sees) is itself a REJECT.
_VERDICT = {
    Disposition.WARN: "ADMIT",
    Disposition.QUARANTINE_REVIEW: "HOLD",
    Disposition.FAIL_SECURE: "REJECT",
}

# Where materialized uploads live inside the bundle. An upload named ``index.*``
# thus lands on the reserved ``uploads/index.md`` and is refused by the path gate.
_MATERIALIZE_PREFIX = "uploads"

# The text formats this slice reads directly (no parser dependency).
_TEXT_SUFFIXES = {".txt", ".md"}

# Zip self-safety caps (OWASP LLM10). Bounded so a decompression bomb is refused
# before its uncompressed bytes are read into memory. Defaults are generous for a
# document inbox; tests pass small caps to exercise the gate.
MAX_ENTRY_BYTES = 5_000_000     # per uncompressed entry
MAX_TOTAL_BYTES = 25_000_000    # per archive, summed across entries


@dataclass(frozen=True)
class Provenance:
    """Where one materialized concept came from — the audit record (brief §6)."""

    concept_path: str
    source_name: str
    source_type: str


@dataclass(frozen=True)
class InboxExtract:
    """The stage-1 output: the OKF bundle, its provenance, and front-end refusals.

    ``rejected`` holds ``(source_name, reason)`` for drops the front-end refuses
    outright (a container threat the guard never gets to see, e.g. a zip bomb).
    Empty for the text-format slice.
    """

    bundle: dict
    provenance: tuple
    rejected: tuple


def _materialize_path(rel_name: str) -> str:
    """Assign an OKF concept path to a dropped file: ``<prefix>/<rel>.md``.

    The relative name is preserved verbatim, including any ``..`` — a zip-slip
    entry (``../../evil.md``) thus lands on a traversal concept path that the
    stage-2 path gate (T4) rejects, rather than being silently normalized away.
    """
    rel = Path(rel_name).with_suffix(".md").as_posix()
    return f"{_MATERIALIZE_PREFIX}/{rel}"


def _is_symlink_entry(info: zipfile.ZipInfo) -> bool:
    """True if a zip entry encodes a Unix symlink (mode bits in external_attr)."""
    return stat.S_ISLNK(info.external_attr >> 16)


def _extract_zip(path, bundle, provenance, rejected, max_entry_bytes, max_total_bytes):
    """Read a ``.zip`` in memory, materializing its text entries; refuse bombs,
    symlinks and oversize entries at the front-end (container threats)."""
    total = 0
    with zipfile.ZipFile(path) as zf:
        for info in zf.infolist():
            name = info.filename
            if name.endswith("/"):
                continue  # directory entry — no content
            if _is_symlink_entry(info):
                rejected.append((name, "symlink entry refused (container threat)"))
                continue
            # Fast reject on the declared uncompressed size (a bomb, before reading).
            if info.file_size > max_entry_bytes:
                rejected.append((name, f"entry exceeds {max_entry_bytes}-byte cap (declared {info.file_size})"))
                continue
            if total + info.file_size > max_total_bytes:
                rejected.append((name, f"archive exceeds {max_total_bytes}-byte total cap"))
                continue
            if Path(name).suffix.lower() not in _TEXT_SUFFIXES:
                continue  # only text concepts are materialized in this slice
            # Bounded read defends against a header that lies about file_size.
            with zf.open(info) as f:
                data = f.read(max_entry_bytes + 1)
            if len(data) > max_entry_bytes:
                rejected.append((name, f"entry expands past {max_entry_bytes}-byte cap on read"))
                continue
            total += len(data)
            concept_path = _materialize_path(name)
            bundle[concept_path] = data.decode("utf-8", errors="replace")
            provenance.append(Provenance(concept_path, name, "zip"))


def extract_inbox(
    paths,
    *,
    max_entry_bytes: int = MAX_ENTRY_BYTES,
    max_total_bytes: int = MAX_TOTAL_BYTES,
) -> InboxExtract:
    """Read dropped files and materialize them into an OKF bundle + provenance.

    ``paths`` is an iterable of file paths. Each ``.txt`` / ``.md`` becomes one
    concept (a ``.md`` keeps its OKF frontmatter verbatim). A ``.zip`` is read in
    memory and its text entries materialized, with bomb/symlink/oversize entries
    refused into ``InboxExtract.rejected``. Other suffixes are not handled yet.
    """
    bundle: dict = {}
    provenance: list = []
    rejected: list = []
    for path in paths:
        path = Path(path)
        suffix = path.suffix.lower()
        if suffix == ".zip":
            _extract_zip(path, bundle, provenance, rejected, max_entry_bytes, max_total_bytes)
        elif suffix in _TEXT_SUFFIXES:
            concept_path = _materialize_path(path.name)
            bundle[concept_path] = path.read_text(encoding="utf-8")
            provenance.append(Provenance(concept_path, path.name, suffix.lstrip(".")))
        else:
            raise ValueError(f"unsupported upload format in this slice: {path.name!r}")
    return InboxExtract(bundle, tuple(provenance), tuple(rejected))


def receive(
    paths,
    *,
    max_entry_bytes: int = MAX_ENTRY_BYTES,
    max_total_bytes: int = MAX_TOTAL_BYTES,
) -> tuple[InboxExtract, BundleResult, str]:
    """The full two-stage inbox: extract & materialize, then guard, then verdict.

    Returns ``(extracted, guard_result, verdict)``. A front-end refusal forces a
    REJECT regardless of the guard's aggregate — the guard never saw that drop.
    """
    extracted = extract_inbox(paths, max_entry_bytes=max_entry_bytes, max_total_bytes=max_total_bytes)
    result = import_bundle(extracted.bundle, origin=Origin.EXTERNAL, channel=Channel.AUTOMATIC)
    verdict = "REJECT" if extracted.rejected else _VERDICT[result.disposition]
    return extracted, result, verdict