rag-security-protocol 0.1.1__py3-none-any.whl

rsp/config.py

@@ -0,0 +1,146 @@
+"""Plugins from a TOML file, so that adding a scanner is not a code change.
+
+Not part of the protocol; a host is free to ignore this module. One that prefers
+another format parses it and passes the mapping on:
+
+    Runtime(plugins_from(yaml.safe_load(text)))
+
+Unknown keys are errors, not warnings: a quietly ignored `on_eror` leaves a
+plugin failing open while the file reads as if it blocks.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import pathlib
+import tomllib
+from typing import Any
+
+from rsp.process import DEFAULT_MAX_OUTPUT, DEFAULT_TIMEOUT
+from rsp.runtime import ConfigError, OnError, Plugin
+
+_TOP_LEVEL = frozenset({"plugins"})
+_SEPARATORS = {"/", os.sep}
+_FIELDS = frozenset({"name", "command", "on_error", "timeout", "max_output"})
+_REQUIRED = ("name", "command")
+
+
+def _text(value: Any, where: str) -> str:
+    if not isinstance(value, str) or not value:
+        raise ConfigError(f"{where}: expected a non-empty string")
+    return value
+
+
+def _command(value: Any, where: str, base: pathlib.Path | None) -> list[str]:
+    """An argv list, never a string.
+
+    A string invites a shell, and splitting it on spaces would make the quoting
+    rules ours to get wrong.
+    """
+    if isinstance(value, str):
+        raise ConfigError(f"{where}: expected a list of arguments, not a string")
+    if not isinstance(value, list) or not value:
+        raise ConfigError(f"{where}: expected a non-empty list of arguments")
+    return [_beside(_text(argument, f"{where}[{at}]"), base) for at, argument in enumerate(value)]
+
+
+def _beside(argument: str, base: pathlib.Path | None) -> str:
+    """A path in the command, read beside the file that names it.
+
+    Resolved against the process instead, a config works from one directory
+    and the user meets that wherever they put theirs. A separator is what
+    tells a path from a name `PATH` resolves, and from an argument that is no
+    path at all.
+    """
+    if base is None or not any(part in argument for part in _SEPARATORS):
+        return argument
+    path = pathlib.PurePath(argument)
+    return argument if path.is_absolute() else str(base / path)
+
+
+def _on_error(value: Any, where: str) -> OnError:
+    if value is None:
+        return OnError.BLOCK
+    try:
+        return OnError(value)
+    except ValueError:
+        allowed = ", ".join(member.value for member in OnError)
+        raise ConfigError(f"{where}: expected one of {allowed}, got {value!r}") from None
+
+
+def _positive(value: Any, where: str) -> float:
+    """A finite positive number: `nan` passes `value <= 0` and would switch off E4."""
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ConfigError(f"{where}: expected a positive number")
+    if not math.isfinite(value) or value <= 0:
+        raise ConfigError(f"{where}: expected a positive finite number, got {value}")
+    return float(value)
+
+
+def _bytes(value: Any, where: str) -> int:
+    """A count of bytes, which a fraction is not: int() would round it away."""
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise ConfigError(f"{where}: expected a whole number of bytes")
+    if value <= 0:
+        raise ConfigError(f"{where}: expected a positive number")
+    return value
+
+
+def _plugin(entry: Any, where: str, base: pathlib.Path | None) -> Plugin:
+    if not isinstance(entry, dict):
+        raise ConfigError(f"{where}: expected a table")
+    if unknown := sorted(set(entry) - _FIELDS):
+        raise ConfigError(f"{where}: unknown key(s) {', '.join(unknown)}")
+    if missing := [field for field in _REQUIRED if field not in entry]:
+        raise ConfigError(f"{where}: missing {', '.join(missing)}")
+
+    return Plugin(
+        name=_text(entry["name"], f"{where}.name"),
+        command=_command(entry["command"], f"{where}.command", base),
+        on_error=_on_error(entry.get("on_error"), f"{where}.on_error"),
+        timeout=_positive(entry.get("timeout", DEFAULT_TIMEOUT), f"{where}.timeout"),
+        max_output=_bytes(entry.get("max_output", DEFAULT_MAX_OUTPUT), f"{where}.max_output"),
+    )
+
+
+def plugins_from(
+    document: Any, where: str = "config", base: pathlib.Path | None = None
+) -> list[Plugin]:
+    """The plugins a mapping describes, in the order it lists them.
+
+    `base` is the directory a relative path in a command is read against,
+    default the process's own. A host parsing some other format passes the
+    directory that format came from.
+    """
+    if not isinstance(document, dict):
+        raise ConfigError(f"{where}: expected a table, got {type(document).__name__}")
+    if unknown := sorted(set(document) - _TOP_LEVEL):
+        raise ConfigError(f"{where}: unknown key(s) {', '.join(unknown)}")
+
+    declared = document.get("plugins", [])
+    if not isinstance(declared, list):
+        raise ConfigError(f"{where}: plugins must be a list of tables")
+    if not declared:
+        # A host with no guards is said out loud, not accepted in silence.
+        raise ConfigError(f"{where}: no plugins configured")
+    plugins = [_plugin(entry, f"{where}: plugins[{at}]", base) for at, entry in enumerate(declared)]
+    # Runtime refuses these too; checking here lets `rsp validate`, which starts nothing, say so.
+    seen: set[str] = set()
+    for plugin in plugins:
+        if plugin.name in seen:
+            raise ConfigError(f"{where}: {plugin.name} configured twice")
+        seen.add(plugin.name)
+    return plugins
+
+
+def load(path: str | pathlib.Path) -> list[Plugin]:
+    """Read a config file. Raises ConfigError for anything it cannot honour."""
+    file = pathlib.Path(path)
+    try:
+        document = tomllib.loads(file.read_text(encoding="utf-8"))
+    except OSError as exc:
+        raise ConfigError(f"{file}: {exc.strerror}") from exc
+    except (tomllib.TOMLDecodeError, UnicodeDecodeError) as exc:
+        raise ConfigError(f"{file}: {exc}") from exc
+    return plugins_from(document, where=str(file), base=file.resolve().parent)

rsp/guards.py

@@ -0,0 +1,91 @@
+"""LlamaIndex adapters: the host side of the protocol.
+
+Pass `store_doc_text=False` to `IngestionPipeline.run` when a docstore is
+configured: it writes documents down a path no transformation sees (K1).
+"""
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.schema import BaseNode, NodeWithScore, QueryBundle, TransformComponent
+from pydantic import ConfigDict
+
+from rsp.runtime import Result, Runtime, Verdict
+
+
+def tag(node: BaseNode, provenance: dict[str, Any]) -> None:
+    """Record findings on the node, out of reach of the embedding and the LLM (Q8).
+
+    LlamaIndex folds metadata into the text it embeds, so a plain update would
+    embed "this chunk held an aws-access-key".
+    """
+    node.metadata.update(provenance)
+    for excluded in (node.excluded_embed_metadata_keys, node.excluded_llm_metadata_keys):
+        excluded.extend(key for key in provenance if key not in excluded)
+
+
+def _judged(runtime: Runtime, node: BaseNode, hook: str, metadata: dict[str, Any]) -> Result:
+    """Judge one node and apply the verdict to it, except for dropping it.
+
+    Dropping is the caller's: one guard reports it and the other shrinks a
+    result set (K1, Q3).
+    """
+    result = runtime.evaluate(hook=hook, content=node.get_content(), metadata=metadata)
+    if not result.blocked:
+        if result.verdict is Verdict.REDACT:
+            # span application is the runtime's job (D6)
+            node.set_content(result.content)
+        tag(node, result.provenance)
+    return result
+
+
+class RSPIngestGuard(TransformComponent):
+    """on_chunk. A blocked node is not returned, so it is never embedded.
+
+    `on_block` is how an operator learns what was dropped: a blocked node leaves
+    no trace in the pipeline's output.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    runtime: Runtime
+    on_block: Callable[[BaseNode, Result], None] | None = None
+
+    def __call__(self, nodes: Sequence[BaseNode], **kwargs: Any) -> Sequence[BaseNode]:
+        kept: list[BaseNode] = []
+        for node in nodes:
+            result = _judged(
+                self.runtime,
+                node,
+                "on_chunk",
+                {"source": node.metadata.get("file_path"), "node_id": node.id_},
+            )
+            if result.blocked:
+                if self.on_block is not None:
+                    self.on_block(node, result)
+                continue
+            kept.append(node)
+        return kept
+
+
+class RSPRetrieveGuard(BaseNodePostprocessor):
+    """on_retrieve. A blocked node is dropped from the result set (Q3)."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    runtime: Runtime
+
+    def _postprocess_nodes(
+        self,
+        nodes: list[NodeWithScore],
+        query_bundle: QueryBundle | None = None,
+    ) -> list[NodeWithScore]:
+        kept: list[NodeWithScore] = []
+        for scored in nodes:
+            metadata = {"node_id": scored.node.id_, "score": scored.score}
+            if _judged(self.runtime, scored.node, "on_retrieve", metadata).blocked:
+                continue
+            kept.append(scored)
+        return kept
+
+
+# No on_response guard: LlamaIndex has no seam for it, and §4 reserves the hook.

rsp/handshake.py

@@ -0,0 +1,77 @@
+"""Handshake: who a plugin says it is, validated (H1, H2)."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from rsp.codec import RSP_VERSION, call
+from rsp.process import DEFAULT_MAX_OUTPUT, DEFAULT_TIMEOUT, Outcome
+
+
+@dataclass(frozen=True)
+class Handshake:
+    """What a plugin says it is (H2).
+
+    Only the fields the host acts on: `hooks` gates dispatch (H3); `version` keys
+    the cache and `deterministic` decides whether caching is legal (D4).
+    """
+
+    rsp_version: str
+    name: str
+    version: str
+    hooks: frozenset[str]
+    deterministic: bool = False
+
+    def supports(self, hook: str) -> bool:
+        return hook in self.hooks
+
+
+def _declaration(payload: Mapping[str, Any]) -> Handshake | None:
+    """Validate a declaration. Unknown fields are ignored, never an error (D8)."""
+    # Absent and wrongly typed are one property, so `get` settles both.
+    hooks = payload.get("hooks")
+    if not isinstance(hooks, list) or not all(isinstance(hook, str) for hook in hooks):
+        return None
+    if not all(isinstance(payload.get(f), str) for f in ("rsp_version", "name", "version")):
+        return None
+
+    # Absent means false: a wrong guess caches an answer that depends on when it was asked.
+    deterministic = payload.get("deterministic", False)
+    if not isinstance(deterministic, bool):
+        return None
+
+    return Handshake(
+        rsp_version=payload["rsp_version"],
+        name=payload["name"],
+        version=payload["version"],
+        hooks=frozenset(hooks),
+        deterministic=deterministic,
+    )
+
+
+def handshake(
+    command: Sequence[str],
+    *,
+    timeout: float = DEFAULT_TIMEOUT,
+    max_output: int = DEFAULT_MAX_OUTPUT,
+) -> tuple[Outcome, Handshake | None]:
+    """Ask a plugin who it is, before sending it any content (H1).
+
+    Its own spawn, since one call is one process (Q6). Caching cannot remove it:
+    the cache key needs the version this call fetches.
+    """
+    reply = call(
+        command,
+        {"rsp_version": RSP_VERSION, "hook": "handshake"},
+        timeout=timeout,
+        max_output=max_output,
+    )
+    if not reply.ok:
+        return reply.outcome, None
+
+    declaration = _declaration(reply.payload)
+    if declaration is None:
+        return Outcome.MALFORMED, None
+    return Outcome.OK, declaration

rsp/ingest.py

@@ -0,0 +1,153 @@
+"""Run a directory through the pipeline and report what did not make it in.
+
+Needs the `llamaindex` extra. The embedding is a stand-in so the demo needs no
+API key: what it shows is which chunks were kept, redacted, or refused.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from bisect import bisect_left
+from functools import cache
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+from rsp.runtime import Plugin, Runtime
+
+if TYPE_CHECKING:
+    from llama_index.core.schema import BaseNode
+
+    from rsp.runtime import Result
+
+# Large enough for a PEM block: a secret split across two chunks is invisible to on_chunk.
+CHUNK = 256
+_READABLE = (".md", ".txt", ".rst")
+
+
+class Finding(NamedTuple):
+    """Where to look, never what was found: quoting a secret makes a second copy (Q8)."""
+
+    verdict: str
+    types: str
+    source: str
+    lines: str
+
+
+class Report(NamedTuple):
+    scanned: int
+    indexed: int
+    findings: list[Finding]
+
+    @property
+    def blocked(self) -> list[Finding]:
+        return [one for one in self.findings if one.verdict == "BLOCK"]
+
+    @property
+    def redacted(self) -> list[Finding]:
+        return [one for one in self.findings if one.verdict == "REDACT"]
+
+
+def documents(directory: pathlib.Path) -> list[Any]:
+    from llama_index.core.schema import Document
+
+    found = sorted(p for p in directory.rglob("*") if p.suffix in _READABLE and p.is_file())
+    if not found:
+        raise FileNotFoundError(f"{directory}: nothing to ingest")
+    return [_document(Document, path) for path in found]
+
+
+def _document(Document: Any, path: pathlib.Path) -> Any:  # noqa: N803
+    """One file, with its path kept out of the text the splitter budgets for.
+
+    The splitter subtracts a document's metadata from the chunk size, so an
+    absolute path in there chunks the same corpus differently depending on
+    where it is checked out — and a scanner then sees different text in one
+    chunk. Excluding it is also what Q8 asks of anything a node carries.
+    """
+    document = Document(text=path.read_text(encoding="utf-8"), metadata={"file_path": str(path)})
+    for excluded in (document.excluded_embed_metadata_keys, document.excluded_llm_metadata_keys):
+        excluded.append("file_path")
+    return document
+
+
+def _newlines(source: str) -> tuple[int, ...]:
+    """Where the line breaks are in a file on disk, read once per version of it."""
+    stat = pathlib.Path(source).stat()
+    return _newlines_as_of(source, stat.st_mtime_ns, stat.st_size)
+
+
+@cache
+def _newlines_as_of(source: str, mtime_ns: int, size: int) -> tuple[int, ...]:
+    """Keyed on when the file was written, so a rewritten one is read again.
+
+    Offsets and not the text: a cache of contents would hold every scanned
+    secret for as long as the process lives.
+    """
+    return tuple(
+        at for at, char in enumerate(pathlib.Path(source).read_text("utf-8")) if char == "\n"
+    )
+
+
+def _lines_of(node: BaseNode) -> str:
+    """The chunk's line range, since the host never sees a span (S4).
+
+    A single line number would send a reader to the top of the chunk as if the
+    secret were there.
+    """
+    source = node.metadata.get("file_path")
+    start = getattr(node, "start_char_idx", None)
+    end = getattr(node, "end_char_idx", None)
+    if source is None or start is None or end is None:
+        return "?"
+    # From the offsets, not the content: a redacted node's text has been rewritten.
+    newlines = _newlines(source)
+    # end is exclusive, so the last line is the one holding its last character.
+    # Taken from end itself, a chunk ending on a newline reaches the empty line
+    # after it; below start, an empty chunk ends before it begins.
+    return f"{bisect_left(newlines, start) + 1}-{bisect_left(newlines, max(start, end - 1)) + 1}"
+
+
+def _finding(verdict: str, provenance: dict[str, Any], node: BaseNode) -> Finding:
+    types = provenance.get("rsp.types") or ["unknown"]
+    return Finding(
+        verdict=verdict,
+        types=", ".join(types),
+        source=node.metadata.get("file_path", "?"),
+        lines=_lines_of(node),
+    )
+
+
+def ingest(directory: pathlib.Path, plugins: list[Plugin], *, store: Any = None) -> Report:
+    """One directory in, a `Report` out; `store` takes the kept chunks, default nowhere.
+
+    A host that wants an index to retrieve from passes one. It is also where
+    the claim is checked: what the index was handed is a stronger answer than
+    a count of what came back.
+    """
+    from llama_index.core.embeddings import MockEmbedding
+    from llama_index.core.ingestion import IngestionPipeline
+    from llama_index.core.node_parser import SentenceSplitter
+
+    from rsp.guards import RSPIngestGuard
+
+    findings: list[Finding] = []
+
+    def record(node: BaseNode, result: Result) -> None:
+        findings.append(_finding("BLOCK", result.provenance, node))
+
+    runtime = Runtime(plugins)  # handshakes before the corpus is read
+    chunks = SentenceSplitter(chunk_size=CHUNK, chunk_overlap=0)(documents(directory))
+    kept = IngestionPipeline(
+        transformations=[
+            RSPIngestGuard(runtime=runtime, on_block=record),
+            MockEmbedding(embed_dim=8),
+        ],
+        # No docstore: it would store documents whole, past every guard (K1).
+        vector_store=store,
+    ).run(nodes=chunks)
+
+    findings += [
+        _finding("REDACT", node.metadata, node)
+        for node in kept
+        if node.metadata.get("rsp.verdict") == "REDACT"
+    ]
+    return Report(scanned=len(chunks), indexed=len(kept), findings=findings)

rsp/process.py

@@ -0,0 +1,183 @@
+"""Subprocess lifecycle: start a plugin, feed it, drain it, reap it.
+
+Bytes in, bytes out, and how the process ended. Knows no JSON, and not what a
+failure means — that is rsp.runtime's job (E1). One process per call (Q6).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import enum
+import io
+import os
+import signal
+import subprocess
+import threading
+import time
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+DEFAULT_TIMEOUT = 5.0  # seconds per call, startup included (E4)
+DEFAULT_MAX_OUTPUT = 1 << 20  # bytes of stdout; more is an error (E1)
+_CHUNK = 65536
+
+
+class Outcome(enum.StrEnum):
+    """How a call ended. Everything but OK is an error (E1, E3).
+
+    One enum for process and content failures alike: the caller only asks
+    whether the response can be trusted.
+    """
+
+    OK = "OK"
+    TIMEOUT = "TIMEOUT"
+    OVERSIZE = "OVERSIZE"
+    UNDELIVERED = "UNDELIVERED"
+    TRUNCATED = "TRUNCATED"
+    CRASHED = "CRASHED"
+    UNSPAWNABLE = "UNSPAWNABLE"
+    EMPTY = "EMPTY"
+    MALFORMED = "MALFORMED"
+    UNENCODABLE = "UNENCODABLE"
+
+
+@dataclass(frozen=True)
+class Invocation:
+    outcome: Outcome
+    stdout: bytes
+    stderr: bytes
+    exit_code: int | None
+    duration: float
+
+    @property
+    def ok(self) -> bool:
+        return self.outcome is Outcome.OK
+
+
+def _kill_group(pgid: int) -> None:
+    """Kill the plugin and everything it spawned.
+
+    A wrapper's tool can outlive it holding the pipe, hence a session of its own
+    at spawn. Takes the group id because os.getpgid() fails once the child is
+    reaped, while the group may still have members.
+    """
+    with contextlib.suppress(OSError):
+        os.killpg(pgid, signal.SIGKILL)
+
+
+def invoke(
+    command: Sequence[str],
+    payload: bytes,
+    *,
+    timeout: float = DEFAULT_TIMEOUT,
+    max_output: int = DEFAULT_MAX_OUTPUT,
+) -> Invocation:
+    """Run one plugin call. Never raises: every failure is an Outcome (E2)."""
+    started = time.monotonic()
+    try:
+        proc = subprocess.Popen(
+            list(command),  # argv, never a shell
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            start_new_session=True,
+        )
+    except (OSError, ValueError, TypeError, IndexError) as exc:
+        # An empty argv raises IndexError and a null entry TypeError (E2).
+        return Invocation(Outcome.UNSPAWNABLE, b"", str(exc).encode(), None, 0.0)
+
+    pgid = proc.pid  # its own group leader; recorded before anything exits
+
+    out, err = bytearray(), bytearray()
+    oversize = threading.Event()
+    undelivered = threading.Event()
+    truncated = threading.Event()
+
+    def feed() -> None:
+        try:
+            proc.stdin.write(payload)
+            proc.stdin.close()
+        except OSError:
+            undelivered.set()  # anything it prints is about content it never saw
+
+    def drain(
+        stream: io.BufferedReader, sink: bytearray, cap: int, *, protocol_channel: bool
+    ) -> None:
+        try:
+            while chunk := stream.read1(_CHUNK):
+                room = cap - len(sink)
+                if len(chunk) <= room:
+                    sink += chunk
+                    continue
+                sink += chunk[:room]
+                if not protocol_channel:
+                    # Noisy is not failing (E3): discard the excess, keep draining.
+                    while stream.read1(_CHUNK):
+                        pass
+                    return
+                oversize.set()
+                _kill_group(pgid)
+                return
+        except OSError:
+            if protocol_channel:  # lost stderr costs only diagnostics (E3)
+                truncated.set()
+
+    workers = [
+        threading.Thread(target=feed, daemon=True),
+        threading.Thread(
+            target=drain,
+            args=(proc.stdout, out, max_output),
+            kwargs={"protocol_channel": True},
+            daemon=True,
+        ),
+        threading.Thread(
+            target=drain,
+            args=(proc.stderr, err, max_output),
+            kwargs={"protocol_channel": False},
+            daemon=True,
+        ),
+    ]
+    for worker in workers:
+        worker.start()
+
+    timed_out = False
+    try:
+        # E4 bounds the call from spawn, so the plugin's startup spends the budget.
+        proc.wait(timeout=max(0.0, timeout - (time.monotonic() - started)))
+    except subprocess.TimeoutExpired:
+        timed_out = True
+        _kill_group(pgid)
+        proc.wait()  # reap: no zombie
+
+    # Also after a clean exit: a wrapper can exit zero with its tool still holding the pipe.
+    _kill_group(pgid)
+
+    for worker in workers:
+        worker.join(timeout=1.0)
+
+    if not any(worker.is_alive() for worker in workers):
+        # Closing a stream a stalled worker still reads is the worse bug.
+        for stream in (proc.stdin, proc.stdout, proc.stderr):
+            with contextlib.suppress(OSError):
+                stream.close()
+
+    if oversize.is_set():
+        outcome = Outcome.OVERSIZE
+    elif timed_out:
+        outcome = Outcome.TIMEOUT
+    elif proc.returncode != 0:
+        outcome = Outcome.CRASHED
+    elif undelivered.is_set():
+        outcome = Outcome.UNDELIVERED
+    elif truncated.is_set():
+        outcome = Outcome.TRUNCATED
+    else:
+        outcome = Outcome.OK
+
+    return Invocation(
+        outcome=outcome,
+        stdout=bytes(out),
+        stderr=bytes(err),
+        exit_code=proc.returncode,
+        duration=time.monotonic() - started,
+    )