"""Persona leakage test suite for Azure AI Search security trimming.

Validates that query paths fail closed and enforce document-level access control.
Tests that no persona retrieves forbidden documents or canary tokens, and verifies
that users are not over-blocked from documents they are authorised to read.
"""

from __future__ import annotations

import contextlib
import re
from itertools import product
from typing import Any

try:
    import pytest
except ImportError:
    # Minimal stand-in allowing the test file to run standalone without pytest installed
    class _DummyMark:
        @staticmethod
        def parametrize(*args: Any, **kwargs: Any) -> Any:
            def decorator(func: Any) -> Any:
                return func
            return decorator

    class _DummyPytest:
        mark = _DummyMark()

        @staticmethod
        @contextlib.contextmanager
        def raises(expected_exc: type[Exception], match: str | None = None) -> Any:
            try:
                yield
            except expected_exc as exc:
                if match and match not in str(exc):
                    raise AssertionError(f"Expected {match!r} in exception message, got {str(exc)!r}") from exc
                return
            raise AssertionError(f"Expected exception {expected_exc.__name__} was not raised")

    pytest = _DummyPytest()  # type: ignore[assignment]


# Synthetic corpus representing documents in an Azure AI Search index.
# group_ids mirrors the Collection(Edm.String) filterable field in Azure AI Search.
# canary is a unique, non-guessable string embedded in the document text.
DOCS: dict[str, dict[str, Any]] = {
    "d1": {
        "group_ids": ["hr"],
        "canary": "CANARY-HR-POLICY-9012",
        "title": "HR Disciplinary Policy",
        "text": "HR disciplinary procedures CANARY-HR-POLICY-9012 confidential.",
    },
    "d2": {
        "group_ids": ["all-staff"],
        "canary": "CANARY-STAFF-HANDBOOK-1104",
        "title": "Employee Handbook",
        "text": "General staff guidelines CANARY-STAFF-HANDBOOK-1104 public to staff.",
    },
    "d3": {
        "group_ids": ["finance"],
        "canary": "CANARY-FIN-BUDGET-4481",
        "title": "Q4 Financial Budget",
        "text": "Q4 revenue forecasts and payroll CANARY-FIN-BUDGET-4481 restricted.",
    },
    "d4": {
        "group_ids": ["hr", "legal"],
        "canary": "CANARY-LEGAL-NOTES-7732",
        "title": "Executive Severance Review",
        "text": "Severance litigation review CANARY-LEGAL-NOTES-7732 hr and legal.",
    },
    "d5": {
        "group_ids": [],
        "canary": "CANARY-ORPHAN-FILE-0000",
        "title": "Unclassified Orphan File",
        "text": "Orphan draft file CANARY-ORPHAN-FILE-0000 with no permissions set.",
    },
}

# Stand-in user directory representing Microsoft Entra ID transitive group memberships.
MEMBERS: dict[str, list[str]] = {
    "u_hr": ["hr", "all-staff"],
    "u_fin": ["finance", "all-staff"],
    "u_guest": [],
    "u_nested": ["legal", "all-staff"],  # Inherited 'legal' through parent group membership
    "u_many": [f"team-{i}" for i in range(250)] + ["hr"],  # Simulates >200 group claim overage
}

# Group IDs must be safe alphanumeric strings, GUIDs, or SharePoint group IDs.
SAFE_GROUP_PATTERN = re.compile(r"^[A-Za-z0-9_:-]+$")


def groups_for(user: str) -> list[str]:
    """Resolve transitive group memberships for a user (simulating Microsoft Graph)."""
    if user == "u_broken":
        raise TimeoutError("Microsoft Graph connection timed out")
    if user not in MEMBERS:
        raise KeyError(f"User {user} not found in directory")
    return MEMBERS[user]


def build_filter(groups: list[str]) -> str | None:
    """Build the OData security filter for Azure AI Search.

    Returns None if the user belongs to no groups, signalling that the index
    must NOT be queried (fail closed).
    """
    if not groups:
        return None
    for g in groups:
        if not SAFE_GROUP_PATTERN.match(g):
            raise ValueError(f"Unsafe group identifier rejected: {g}")
    return "group_ids/any(g:search.in(g, '" + ",".join(groups) + "'))"


def execute_search(user: str, query: str, enforce_filter: bool = True) -> dict[str, Any]:
    """Stand-in query execution simulating Azure AI Search retrieval."""
    try:
        groups = groups_for(user)
        security_filter = build_filter(groups)
    except Exception:
        # Fail closed: on any directory error or missing user, return empty results
        return {"ids": [], "text": "", "count": 0}

    # If the user has no groups, fail closed without invoking the index
    if enforce_filter and security_filter is None:
        return {"ids": [], "text": "", "count": 0}

    if not enforce_filter:
        # Deliberate vulnerability path: skips the filter
        matched_ids = list(DOCS.keys())
    else:
        # Trimming logic: document is visible only if group_ids overlaps with user groups
        user_group_set = set(groups)
        matched_ids = [
            doc_id
            for doc_id, doc in DOCS.items()
            if set(doc["group_ids"]) & user_group_set
        ]

    returned_text = " ".join(DOCS[doc_id]["text"] for doc_id in matched_ids)
    return {
        "ids": matched_ids,
        "text": returned_text,
        "count": len(matched_ids),
    }


# Standard application query paths under test
PATHS = {
    "search": lambda user, q: execute_search(user, q, enforce_filter=True),
    "chat": lambda user, q: execute_search(user, q, enforce_filter=True),
    "suggest": lambda user, q: execute_search(user, q, enforce_filter=True),
    "agent": lambda user, q: execute_search(user, q, enforce_filter=True),
    "open": lambda user, q: execute_search(user, q, enforce_filter=True),
    "export": lambda user, q: execute_search(user, q, enforce_filter=True),
}

ALL_DOC_IDS = set(DOCS.keys())

# Test personas with explicit allow (must see) and forbid (must never see) sets
PERSONAS = [
    {
        "name": "hr_member",
        "user": "u_hr",
        "allow": {"d1", "d2", "d4"},
        "forbid": {"d3", "d5"},
    },
    {
        "name": "finance_member",
        "user": "u_fin",
        "allow": {"d2", "d3"},
        "forbid": {"d1", "d4", "d5"},
    },
    {
        "name": "guest_unmapped",
        "user": "u_guest",
        "allow": set(),
        "forbid": ALL_DOC_IDS,
    },
    {
        "name": "nested_group_member",
        "user": "u_nested",
        "allow": {"d2", "d4"},
        "forbid": {"d1", "d3", "d5"},
    },
    {
        "name": "overage_many_groups",
        "user": "u_many",
        "allow": {"d1", "d4"},
        "forbid": {"d2", "d3", "d5"},
    },
    {
        "name": "directory_lookup_timeout",
        "user": "u_broken",
        "allow": set(),
        "forbid": ALL_DOC_IDS,
    },
    {
        "name": "deleted_user_account",
        "user": "u_gone",
        "allow": set(),
        "forbid": ALL_DOC_IDS,
    },
]


@pytest.mark.parametrize(
    "path_name,persona",
    list(product(PATHS.keys(), PERSONAS)),
    ids=lambda x: x if isinstance(x, str) else x["name"],
)
def test_path_enforces_security_trimming(path_name: str, persona: dict[str, Any]) -> None:
    """Verify that every query path returns only authorised documents for each persona."""
    path_func = PATHS[path_name]
    result = path_func(persona["user"], "policy guidelines")
    returned_ids = set(result["ids"])

    # 1. Leakage check: forbidden document IDs must never appear in results
    leaked_ids = returned_ids & persona["forbid"]
    assert not leaked_ids, (
        f"Security leak on path '{path_name}': persona '{persona['name']}' "
        f"retrieved forbidden document(s): {sorted(leaked_ids)}"
    )

    # 2. Canary check: canary tokens of forbidden documents must never appear in returned text
    for forbidden_id in persona["forbid"]:
        canary = DOCS[forbidden_id]["canary"]
        assert canary not in result["text"], (
            f"Content leak on path '{path_name}': persona '{persona['name']}' "
            f"received canary token '{canary}' from restricted document '{forbidden_id}'"
        )

    # 3. Over-blocking check: all allowed documents must be returned
    missing_allowed = persona["allow"] - returned_ids
    assert not missing_allowed, (
        f"Over-blocking on path '{path_name}': persona '{persona['name']}' "
        f"could not access permitted document(s): {sorted(missing_allowed)}"
    )


def test_build_filter_normal() -> None:
    """Validate OData filter construction with valid group lists."""
    flt = build_filter(["hr", "finance"])
    assert flt == "group_ids/any(g:search.in(g, 'hr,finance'))"


def test_build_filter_empty_fails_closed() -> None:
    """Validate that an empty group list returns None (do not call search index)."""
    assert build_filter([]) is None


@pytest.mark.parametrize(
    "hostile_input",
    [
        "hr' or '1'='1",
        "hr,finance; drop index",
        "team 1",
        "legal|admin",
        "hr' or true or ('",
        "",
    ],
)
def test_build_filter_rejects_unsafe_group_identifiers(hostile_input: str) -> None:
    """Ensure injection attempts in group identifiers are rejected before OData generation."""
    with pytest.raises(ValueError, match="Unsafe group identifier rejected"):
        build_filter([hostile_input])


def test_deliberate_leak_is_caught() -> None:
    """Verify that an unprotected path skipping filters triggers test assertion failure."""
    insecure_export = lambda user, q: execute_search(user, q, enforce_filter=False)
    target_persona = PERSONAS[0]  # hr_member: forbidden d3 and d5
    result = insecure_export(target_persona["user"], "policy")
    returned_ids = set(result["ids"])

    # Verify that the test correctly catches the deliberate leak
    assert returned_ids & target_persona["forbid"] == {"d3", "d5"}
    assert DOCS["d3"]["canary"] in result["text"]


if __name__ == "__main__":
    count = 0
    for path_name, persona in product(PATHS.keys(), PERSONAS):
        test_path_enforces_security_trimming(path_name, persona)
        count += 1
    test_build_filter_normal()
    test_build_filter_empty_fails_closed()
    for bad in ["hr' or '1'='1", "hr,finance; drop index", "team 1", "legal|admin", "hr' or true or ('", ""]:
        test_build_filter_rejects_unsafe_group_identifiers(bad)
    test_deliberate_leak_is_caught()
    print(f"All {count} persona path test combinations and filter validation checks passed successfully.")