leocrm/app/ai/data_policy.py

"""Runtime provider / data policy enforcement for AI agents.

Filters messages and context before they reach the LLM based on:
- Sensitive fields (``app.core.sensitive_data.SENSITIVE_FIELDS``)
- AI use-case metadata (allowed data categories)
- Provider compliance (data residency / allowed data classes)

This is the enforcement layer that guarantees an agent never sends data it
is not permitted to process to a provider that is not approved for it.
"""

from __future__ import annotations

import logging
import uuid
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.ai.ai_use_case import AIUseCaseMetadata
from app.core.sensitive_data import (
    SENSITIVE_FIELDS,
)

logger = logging.getLogger(__name__)

# Data categories that map to entity types for sensitive-field filtering.
_CATEGORY_ENTITY_MAP = {
    "contact_data": "contact",
    "email_content": "mail_account",
    "communication": "mail_account",
}


async def enforce_data_policy(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    messages: list[dict[str, Any]],
    agent_definition: Any,
) -> list[dict[str, Any]]:
    """Filter messages/context based on the data policy.

    Steps:
    1. Remove sensitive fields from any dict content in the messages.
    2. Check AI use-case metadata for allowed data categories.
    3. Check provider compliance for data residency requirements.

    Args:
        db: Async DB session (may be ``None`` in tests / mock mode).
        tenant_id: Tenant ID for provider lookup.
        messages: The chat messages to filter.
        agent_definition: AgentDefinition with ``ai_use_case_metadata``.

    Returns:
        A new list of messages with disallowed data removed.
    """
    metadata = AIUseCaseMetadata.from_dict(
        getattr(agent_definition, "ai_use_case_metadata", None)
    )

    # Provider compliance (data residency / allowed data classes).
    compliance: dict[str, Any] | None = None
    if db is not None and tenant_id is not None:
        try:
            from app.ai.llm_client import get_provider_compliance

            compliance = await get_provider_compliance(db, tenant_id)
        except Exception:
            logger.debug("Failed to load provider compliance — skipping residency check")

    filtered: list[dict[str, Any]] = []
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, dict):
            content = _filter_dict_content(
                content, metadata, compliance, agent_definition
            )
        elif isinstance(content, list):
            content = [
                _filter_dict_content(c, metadata, compliance, agent_definition)
                if isinstance(c, dict)
                else c
                for c in content
            ]
        elif isinstance(content, str):
            # F14 (Astra P1): JSON-serialized strings passed through
            # UNFILTERED before — a payload like '{"smtp_password": ...}'
            # reached the provider verbatim. Parse, filter, re-serialize;
            # non-JSON strings stay unchanged (plain prose is fine).
            content = _filter_json_string_content(
                content, metadata, compliance, agent_definition
            )
        new_msg = dict(msg)
        new_msg["content"] = content
        filtered.append(new_msg)

    return filtered


def _filter_json_string_content(
    content: str,
    metadata: AIUseCaseMetadata,
    compliance: dict[str, Any] | None,
    agent_definition: Any,
) -> str:
    """Filter a JSON-serialized string payload (F14).

    Tries to parse the string as a JSON object/array and runs the SAME
    dict-level filtering on it. Non-JSON strings are returned unchanged.
    """
    stripped = content.strip()
    if not stripped.startswith(("{", "[")):
        return content
    try:
        import json

        parsed = json.loads(stripped)
    except (json.JSONDecodeError, ValueError):
        return content  # not JSON — plain string content is not a leak vector
    if isinstance(parsed, dict):
        filtered = _filter_dict_content(parsed, metadata, compliance, agent_definition)
        import json

        return json.dumps(filtered)
    if isinstance(parsed, list):
        filtered = [
            _filter_dict_content(item, metadata, compliance, agent_definition)
            if isinstance(item, dict)
            else item
            for item in parsed
        ]
        import json

        return json.dumps(filtered)
    return content


def _filter_dict_content(
    data: dict[str, Any],
    metadata: AIUseCaseMetadata,
    compliance: dict[str, Any] | None,
    agent_definition: Any,
) -> dict[str, Any]:
    """Filter a single dict (entity payload) against the data policy."""
    # 1. Remove sensitive fields (always blocked from LLM context).
    result = _strip_sensitive_fields(data)

    # 2. Enforce allowed data categories from AI use-case metadata.
    if metadata.data_categories:
        result = _filter_by_allowed_categories(result, metadata.data_categories)

    # 3. Provider compliance — block fields whose data class the provider
    #    is not approved to process.
    if compliance is not None:
        result = _filter_by_provider_compliance(result, compliance)

    return result


def _strip_sensitive_fields(data: dict[str, Any]) -> dict[str, Any]:
    """Recursively remove any key that matches a sensitive field name."""
    sensitive_names = set()
    for fields in SENSITIVE_FIELDS.values():
        sensitive_names |= fields

    result: dict[str, Any] = {}
    for key, value in data.items():
        if key in sensitive_names:
            continue
        if isinstance(value, dict):
            result[key] = _strip_sensitive_fields(value)
        elif isinstance(value, list):
            result[key] = [
                _strip_sensitive_fields(v) if isinstance(v, dict) else v
                for v in value
            ]
        else:
            result[key] = value
    return result


def _filter_by_allowed_categories(
    data: dict[str, Any], allowed_categories: list[str]
) -> dict[str, Any]:
    """Remove entity-type payloads whose category is not allowed.

    Uses the category→entity mapping to decide whether a dict represents a
    disallowed entity type. Unknown dicts are kept (fail-open for generic
    context that has no clear entity type).
    """
    # Determine the entity type of this dict by checking for known keys.
    entity_type = _guess_entity_type(data)
    if entity_type is None:
        return data

    category = _entity_to_category(entity_type)
    if category is not None and category not in allowed_categories:
        return {}
    return data


def _filter_by_provider_compliance(
    data: dict[str, Any], compliance: dict[str, Any]
) -> dict[str, Any]:
    """Remove fields whose data class the provider may not process."""
    allowed_classes = compliance.get("allowed_data_classes") or []
    if not allowed_classes:
        return data  # No restriction configured (fail-open).

    from app.core.sensitive_data import check_provider_compliance

    result: dict[str, Any] = {}
    for key, value in data.items():
        if isinstance(value, dict):
            result[key] = _filter_by_provider_compliance(value, compliance)
            continue
        # Determine data class for this field (best-effort).
        data_class = _guess_data_class(key, value)
        if check_provider_compliance(allowed_classes, data_class):
            result[key] = value
    return result


def _guess_entity_type(data: dict[str, Any]) -> str | None:
    """Best-effort guess of the entity type from dict keys."""
    if any(k in data for k in ("email", "smtp_password", "imap_password")):
        return "mail_account"
    if any(k in data for k in ("first_name", "last_name", "company_id")):
        return "contact"
    if any(k in data for k in ("secret_key", "encryption_key")):
        return "system_settings"
    return None


def _entity_to_category(entity_type: str) -> str | None:
    """Map an entity type to a data category."""
    for category, entity in _CATEGORY_ENTITY_MAP.items():
        if entity == entity_type:
            return category
    return None


def _guess_data_class(key: str, value: Any) -> str:
    """Best-effort data class for a field (defaults to 'internal')."""
    # Sensitive field names are always critical.
    for fields in SENSITIVE_FIELDS.values():
        if key in fields:
            return "critical"
    # Heuristic: values that look like credentials/tokens are critical.
    if isinstance(value, str) and any(
        marker in key.lower() for marker in ("password", "token", "secret", "key")
    ):
        return "critical"
    return "internal"