"""Shared secure retrieval function for Azure AI Search security-filter trimming (PB-RAG-04-02). Every query path (search, chat grounding, agent tools, suggestions) calls this module. No other code may call the index directly. Rules this module enforces: - The caller passes the Entra ID object id (oid) of a user the application has already authenticated. The filter only trims results; it does not authenticate anyone. - Group membership comes from Microsoft Graph transitiveMemberOf, so nested groups are honoured. The token's groups claim is not used, because it can be missing (overage) or stale. - The security filter uses search.in and is always ANDed with business filters, never ORed. - Business filters come from an allow-list of fields with escaped values, never from raw strings. - Any identity, directory or search error returns no results (fail closed). - The audit log records the user id, group count and a filter hash, never the group list. Environment: SEARCH_ENDPOINT, SEARCH_INDEX (index or alias name), GROUP_CACHE_TTL_SECONDS, ALLOWED_FILTER_FIELDS (comma separated), VECTOR_FIELD (optional). """ from __future__ import annotations import hashlib import logging import os import re import threading import time from dataclasses import dataclass from typing import Any, Callable, Mapping, Sequence import requests API_VERSION = "2026-04-01" GRAPH_BASE = "https://graph.microsoft.com/v1.0" GRAPH_SCOPE = "https://graph.microsoft.com/.default" SEARCH_SCOPE = "https://search.azure.com/.default" GROUP_TYPE = "#microsoft.graph.group" GUID = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$") FIELD_NAME = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$") log = logging.getLogger("secure_search") TokenProvider = Callable[[str], str] # scope -> bearer token class TrimmingError(Exception): """Raised when the caller's groups cannot be resolved. Callers never see it: they get no results.""" @dataclass(frozen=True) class Settings: search_endpoint: str index_name: str cache_ttl_seconds: int allowed_filter_fields: frozenset[str] vector_field: str = "content_vector" timeout_seconds: float = 10.0 max_graph_pages: int = 100 @classmethod def from_env(cls) -> "Settings": return cls( search_endpoint=os.environ["SEARCH_ENDPOINT"].rstrip("/"), index_name=os.environ["SEARCH_INDEX"], # No default on purpose: the TTL is the revocation delay agreed in the PB-RAG-04-01 decision record. cache_ttl_seconds=int(os.environ["GROUP_CACHE_TTL_SECONDS"]), allowed_filter_fields=frozenset( f.strip() for f in os.environ.get("ALLOWED_FILTER_FIELDS", "").split(",") if f.strip() ), vector_field=os.environ.get("VECTOR_FIELD", "content_vector"), ) def default_token_provider() -> TokenProvider: """Managed identity in Azure, developer sign-in locally.""" from azure.identity import DefaultAzureCredential credential = DefaultAzureCredential() return lambda scope: credential.get_token(scope).token def _normalise_guid(value: str, what: str) -> str: v = (value or "").strip().lower() if not GUID.match(v): raise ValueError(f"{what} is not an Entra ID object id") return v class GroupCache: """Per-process cache of transitive group ids. Keep the TTL at or below the agreed revocation delay.""" def __init__(self, ttl_seconds: int, clock: Callable[[], float] = time.monotonic) -> None: self._ttl = ttl_seconds self._clock = clock self._items: dict[str, tuple[float, list[str]]] = {} self._lock = threading.Lock() def get_or_load(self, user_oid: str, loader: Callable[[str], list[str]]) -> list[str]: now = self._clock() with self._lock: hit = self._items.get(user_oid) if hit and hit[0] > now: return hit[1] groups = loader(user_oid) # errors propagate and are not cached with self._lock: self._items[user_oid] = (now + self._ttl, groups) return groups def evict(self, user_oid: str) -> None: """Call on sign-out or on a revocation event (PB-RAG-04-03).""" with self._lock: self._items.pop(user_oid, None) def resolve_groups(user_oid: str, token_provider: TokenProvider, session: Any, settings: Settings) -> list[str]: """Return the user's direct and transitive group ids from Microsoft Graph. No OData cast or $filter: those need ConsistencyLevel: eventual, which reads an index that might not be up to date. Non-group objects (directory roles, administrative units) are dropped here. """ oid = _normalise_guid(user_oid, "user_oid") headers = {"Authorization": f"Bearer {token_provider(GRAPH_SCOPE)}"} url: str | None = f"{GRAPH_BASE}/users/{oid}/transitiveMemberOf" groups: set[str] = set() pages = 0 while url: pages += 1 if pages > settings.max_graph_pages: raise TrimmingError("too many membership pages") resp = session.get(url, headers=headers, timeout=settings.timeout_seconds) if resp.status_code == 429: raise TrimmingError(f"Graph throttled, Retry-After={resp.headers.get('Retry-After')}") if resp.status_code != 200: raise TrimmingError(f"Graph returned {resp.status_code}") body = resp.json() for item in body.get("value", []): if item.get("@odata.type") == GROUP_TYPE: groups.add(_normalise_guid(item.get("id", ""), "group id")) url = body.get("@odata.nextLink") return sorted(groups) def build_security_filter(user_oid: str, group_ids: Sequence[str]) -> str: """Documents shared with the user directly, or with any group the user belongs to.""" oid = _normalise_guid(user_oid, "user_oid") clauses = [f"user_ids/any(u: u eq '{oid}')"] clean = [_normalise_guid(g, "group id") for g in group_ids] if clean: clauses.append(f"group_ids/any(g: search.in(g, '{','.join(clean)}', ','))") return "(" + " or ".join(clauses) + ")" def _literal(value: Any) -> str: if isinstance(value, bool): return "true" if value else "false" if isinstance(value, (int, float)): return repr(value) return "'" + str(value).replace("'", "''") + "'" # OData escapes a quote by doubling it def build_business_filter(filters: Mapping[str, Any] | None, allowed: frozenset[str]) -> str | None: """Turn {field: value | [values]} into an OData filter. Raw filter strings are not accepted.""" if not filters: return None clauses = [] for name, value in filters.items(): if name not in allowed or not FIELD_NAME.match(name): raise ValueError(f"filter field not allowed: {name!r}") if name in ("user_ids", "group_ids"): raise ValueError("security fields cannot be used as business filters") if isinstance(value, (list, tuple, set)): values = [str(v) for v in value] if not values or any("|" in v for v in values): raise ValueError(f"invalid value list for {name!r}") joined = "|".join(v.replace("'", "''") for v in values) clauses.append(f"search.in({name}, '{joined}', '|')") else: clauses.append(f"{name} eq {_literal(value)}") return " and ".join(f"({c})" for c in clauses) def _filter_hash(flt: str) -> str: return hashlib.sha256(flt.encode("utf-8")).hexdigest()[:16] class SecureRetriever: """The only object allowed to query the index. Create one per process and share it.""" def __init__(self, settings: Settings, token_provider: TokenProvider | None = None, session: Any = None, cache: GroupCache | None = None) -> None: self.settings = settings self.token_provider = token_provider or default_token_provider() self.session = session or requests.Session() self.cache = cache or GroupCache(settings.cache_ttl_seconds) def _filter_for(self, user_oid: str, business_filters: Mapping[str, Any] | None) -> str | None: """Return the full filter, or None when the caller must get no results.""" try: groups = self.cache.get_or_load( _normalise_guid(user_oid, "user_oid"), lambda oid: resolve_groups(oid, self.token_provider, self.session, self.settings), ) security = build_security_filter(user_oid, groups) except Exception as exc: # noqa: BLE001 - any failure must fail closed log.warning("trim_fail_closed reason=%s", type(exc).__name__) return None business = build_business_filter(business_filters, self.settings.allowed_filter_fields) flt = f"{security} and ({business})" if business else security log.info("trim_applied user=%s groups=%d filter_sha256=%s", user_oid, len(groups), _filter_hash(flt)) return flt def _post(self, operation: str, body: dict[str, Any]) -> dict[str, Any] | None: url = (f"{self.settings.search_endpoint}/indexes/{self.settings.index_name}" f"/docs/{operation}?api-version={API_VERSION}") headers = {"Authorization": f"Bearer {self.token_provider(SEARCH_SCOPE)}"} try: resp = self.session.post(url, headers=headers, json=body, timeout=self.settings.timeout_seconds) except Exception as exc: # noqa: BLE001 log.warning("search_failed operation=%s reason=%s", operation, type(exc).__name__) return None if resp.status_code != 200: log.warning("search_failed operation=%s status=%s", operation, resp.status_code) return None return resp.json() def search(self, query: str, user_oid: str, *, vector: Sequence[float] | None = None, business_filters: Mapping[str, Any] | None = None, top: int = 10, select: Sequence[str] | None = None) -> list[dict[str, Any]]: """Keyword, vector or hybrid search, always trimmed. Returns [] when trimming cannot be applied.""" flt = self._filter_for(user_oid, business_filters) if flt is None: return [] body: dict[str, Any] = {"search": query or "*", "filter": flt, "top": top} if select: body["select"] = ",".join(select) if vector is not None: # preFilter keeps the trim inside the HNSW traversal. Never add vectorQueries.filterOverride: # it replaces the top-level filter, including this security filter. body["vectorQueries"] = [{"kind": "vector", "vector": list(vector), "fields": self.settings.vector_field, "k": top}] body["vectorFilterMode"] = "preFilter" result = self._post("search", body) hits = result.get("value", []) if result else [] log.info("trim_result user=%s count=%d", user_oid, len(hits)) return hits def suggest(self, prefix: str, user_oid: str, suggester_name: str, top: int = 5) -> list[dict[str, Any]]: """Suggestions go through the same filter; an untrimmed suggest call leaks titles.""" flt = self._filter_for(user_oid, None) if flt is None: return [] body = {"search": prefix, "suggesterName": suggester_name, "filter": flt, "top": top} result = self._post("suggest", body) return result.get("value", []) if result else []