diff --git a/src/portfolio_optimiser/retrieval.py b/src/portfolio_optimiser/retrieval.py new file mode 100644 index 0000000..249605f --- /dev/null +++ b/src/portfolio_optimiser/retrieval.py @@ -0,0 +1,130 @@ +"""Framework-agnostic local-folder retriever — the D7-portable data seam. + +**NO** ``agent_framework``, **NO** ``mcp`` — pure stdlib. Documents are chunked at ingest +with EXACT character spans, so every retrieved chunk carries a locator that slices its own +snippet byte-for-byte (citations exact by construction — research 02). Scoring is +keyword/substring overlap (no heavy vector deps — D5/D6). + +Path-security (this module reads a docs folder; the EscapeRoute CVE class, research 02 +Dim 5): every accessed path is canonicalised with ``os.path.realpath`` and boundary-checked +against the canonicalised ``docs_dir``. A ``..`` traversal, a symlink escaping the folder, +and a sibling directory sharing a name prefix (``/a/docs`` vs ``/a/docs-evil``) are all +rejected fail-closed. Native file APIs only. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path + +_CHUNK_CHARS = 240 # paragraph-ish window; small + deterministic for the synthetic corpus + + +class PathSecurityError(RuntimeError): + """A requested path escapes ``docs_dir`` (traversal / symlink / prefix-collision).""" + + +@dataclass(frozen=True) +class TextSpan: + """An exact half-open character span ``[start_index, end_index)`` into a document's + text. The SINGLE owner of this type — ``provenance.Citation.locator`` imports it from + here (Step 6).""" + + start_index: int + end_index: int + + +@dataclass(frozen=True) +class RetrievedChunk: + """One retrieved chunk with an exact-by-construction citation locator.""" + + file: str + locator: TextSpan + snippet: str + score: float + + +def is_within_dir(candidate: str, docs_dir: str) -> bool: + """True iff the realpath of ``candidate`` is inside the realpath of ``docs_dir``. + + Fail-closed against symlink-escape and prefix-collision: uses ``os.path.commonpath`` on + canonical paths, so a sibling dir sharing a name prefix (``/a/docs-evil`` vs ``/a/docs``) + is correctly rejected (a naive ``startswith`` would not).""" + real_root = os.path.realpath(docs_dir) + real_candidate = os.path.realpath(candidate) + try: + return os.path.commonpath([real_root, real_candidate]) == real_root + except ValueError: + # Different drives / mixed absolute-relative -> not within. + return False + + +def safe_resolve(docs_dir: str, relative: str) -> str: + """Resolve ``relative`` against ``docs_dir`` and return its canonical path, or raise + ``PathSecurityError`` if it escapes the folder (``..`` traversal, symlink escape, or a + prefix-collision sibling). Fail-closed.""" + target = os.path.join(os.path.realpath(docs_dir), relative) + if not is_within_dir(target, docs_dir): + raise PathSecurityError(f"path escapes docs_dir: {relative!r}") + return os.path.realpath(target) + + +def _resolve_docs_dir(docs_dir: str) -> str: + real = os.path.realpath(docs_dir) + if not os.path.isdir(real): + raise ValueError(f"docs_dir is not a directory: {docs_dir}") + return real + + +def _chunk_text(text: str) -> list[TextSpan]: + """Split text into fixed-size character windows, each an exact half-open span.""" + spans: list[TextSpan] = [] + for start in range(0, len(text), _CHUNK_CHARS): + end = min(start + _CHUNK_CHARS, len(text)) + if end > start: + spans.append(TextSpan(start_index=start, end_index=end)) + return spans + + +def _score(query: str, snippet: str) -> float: + """Keyword overlap: fraction of distinct query tokens present in the snippet (lowered). + Deterministic, no vectors.""" + q_tokens = {t for t in query.lower().split() if t} + if not q_tokens: + return 0.0 + blob = snippet.lower() + hits = sum(1 for t in q_tokens if t in blob) + return hits / len(q_tokens) + + +def retrieve(query: str, docs_dir: str, top_k: int = 3) -> list[RetrievedChunk]: + """Return the top-``top_k`` chunks across the docs folder ranked by keyword overlap. + + Every file is boundary-checked against ``docs_dir`` before it is read (fail-closed: a + symlink whose realpath escapes the folder is skipped, never opened). Locators are exact + by construction: ``snippet == file_text[start:end]``. Deterministic: ties break by + ``(file, start_index)``.""" + if top_k <= 0: + raise ValueError(f"top_k must be positive, got {top_k}") + root = _resolve_docs_dir(docs_dir) + candidates: list[RetrievedChunk] = [] + for entry in sorted(Path(root).rglob("*")): + if not entry.is_file(): + continue + if not is_within_dir(str(entry), root): + continue # symlink escaping the folder -> fail closed, never read + try: + text = entry.read_text(encoding="utf-8") + except UnicodeDecodeError: + continue # skip binary + rel = os.path.relpath(os.path.realpath(str(entry)), root) + for span in _chunk_text(text): + snippet = text[span.start_index : span.end_index] + candidates.append( + RetrievedChunk( + file=rel, locator=span, snippet=snippet, score=_score(query, snippet) + ) + ) + ranked = sorted(candidates, key=lambda c: (-c.score, c.file, c.locator.start_index)) + return ranked[:top_k] diff --git a/tests/test_retrieval.py b/tests/test_retrieval.py new file mode 100644 index 0000000..63a4a71 --- /dev/null +++ b/tests/test_retrieval.py @@ -0,0 +1,86 @@ +"""Step 5 tests — local-folder retriever core (exact locators + EscapeRoute path-security). + +Citations are exact by construction (locator slices the snippet byte-for-byte) and +deterministic; the path-security checks are behavioural (traversal / symlink-escape / +prefix-collision sibling all fail-closed). Pattern: tests/test_reference_domain.py + +tests/test_backends.py (raises). +""" + +import os + +import pytest + +from portfolio_optimiser.retrieval import ( + PathSecurityError, + is_within_dir, + retrieve, + safe_resolve, +) + + +@pytest.fixture() +def docs(tmp_path): + d = tmp_path / "docs" + d.mkdir() + (d / "asphalt.txt").write_text( + "Asphalt Ab11 unit rate renegotiation reduced the paving cost on the school stretch.", + encoding="utf-8", + ) + (d / "drainage.txt").write_text( + "Drainage length was reduced after a survey on the low-traffic section.", + encoding="utf-8", + ) + return d + + +def test_locator_exactly_slices_snippet(docs) -> None: + hits = retrieve("asphalt paving cost", str(docs), top_k=5) + assert hits + for h in hits: + text = (docs / h.file).read_text(encoding="utf-8") + assert text[h.locator.start_index : h.locator.end_index] == h.snippet + + +def test_retrieve_is_deterministic(docs) -> None: + a = retrieve("drainage survey", str(docs), top_k=3) + b = retrieve("drainage survey", str(docs), top_k=3) + assert [(h.file, h.locator) for h in a] == [(h.file, h.locator) for h in b] + + +def test_top_k_respected(docs) -> None: + assert len(retrieve("the", str(docs), top_k=1)) <= 1 + + +def test_non_positive_top_k_rejected(docs) -> None: + with pytest.raises(ValueError): + retrieve("x", str(docs), top_k=0) + + +def test_parent_traversal_rejected(docs) -> None: + # A `..` traversal escaping the docs folder is rejected fail-closed. + with pytest.raises(PathSecurityError): + safe_resolve(str(docs), "../outside.txt") + + +def test_symlink_escape_is_not_read(tmp_path, docs) -> None: + # A symlink INSIDE docs_dir pointing OUTSIDE must never be read (fail-closed). + secret = tmp_path / "secret.txt" + secret.write_text("TOPSECRET asphalt asphalt asphalt exfiltrated", encoding="utf-8") + os.symlink(str(secret), str(docs / "link.txt")) + # safe_resolve refuses the escaping symlink... + with pytest.raises(PathSecurityError): + safe_resolve(str(docs), "link.txt") + # ...and retrieve never surfaces the outside content even on a matching query. + hits = retrieve("asphalt", str(docs), top_k=10) + assert all("TOPSECRET" not in h.snippet for h in hits) + + +def test_prefix_collision_sibling_rejected(tmp_path) -> None: + # A sibling dir sharing a name prefix (docs vs docs-evil) is NOT inside docs. + root = tmp_path / "docs" + root.mkdir() + sibling = tmp_path / "docs-evil" + sibling.mkdir() + intruder = sibling / "secret.txt" + intruder.write_text("x", encoding="utf-8") + assert is_within_dir(str(intruder), str(root)) is False