obsidian-wiki/obsidian_wiki/mcp_assets.py

"""MCP-facing source file lookup and access control helpers."""

from __future__ import annotations

import json
import mimetypes
import os
import re
from pathlib import Path
from typing import Any

from obsidian_wiki.access_control import as_list, load_profiles, source_decision
from obsidian_wiki.config import config_value
from obsidian_wiki.file_assets import (
    absolute_asset_path,
    clean_config_value,
    db_path_for_vault,
    get_asset_by_file_id,
    registry_path_for_vault,
    search_assets,
)

DEFAULT_PROFILE_NAME = "bjd-general-bot"
DEFAULT_MAX_FILE_BYTES = 25 * 1024 * 1024

ALLOWED_ENGINEERING_DOC_TYPES = {
    "drawing",
    "customer_3d_model",
    "test_report",
    "product_parameter",
    "customer_specification",
    "product_specification",
}

BLOCKED_DOC_TYPES = {
    "order_data",
    "customer_product_price",
    "customer_order_detail",
    "delivery_detail",
    "email_correspondence",
    "evidence_media",
    "claim_record",
    "customer_complaint_source",
    "customer_complaint_report",
    "inspection_control",
    "personal_document",
    "identity_document",
}

BLOCKED_CATEGORY_PREFIXES = (
    "orders",
    "orders/",
    "customer-pricing",
    "customer-order-detail",
    "pricing-detail",
    "quality/customer-complaints",
    "quality/claims",
    "claims",
    "emails",
    "email",
    "pii",
    "personal",
)

SAFE_RESULT_FIELDS = (
    "file_id",
    "resource_uri",
    "score",
    "original_name",
    "title",
    "summary",
    "doc_type",
    "source_type",
    "mime_type",
    "size_bytes",
    "project",
    "visibility",
    "data_classification",
    "product_lines",
    "customers",
    "competitors",
    "standards",
    "business_categories",
    "bjd_part_numbers",
    "competitor_part_numbers",
    "series",
    "pitch",
    "pin_count",
    "related_pages",
    "classification_confidence",
    "updated_at",
)

ASSET_URI_RE = re.compile(r"^asset://source/sha256/([0-9a-fA-F]{64})$")


class AssetAccessError(RuntimeError):
    """Raised when ACL denies access to a file asset."""


class AssetNotFoundError(FileNotFoundError):
    """Raised when an indexed asset or archived file is missing."""


class AssetTooLargeError(RuntimeError):
    """Raised when a resource exceeds the configured read limit."""


def asset_uri_for_file_id(file_id: str) -> str:
    digest = file_id.removeprefix("sha256:")
    if not re.fullmatch(r"[0-9a-fA-F]{64}", digest):
        raise ValueError(f"unsupported file_id for MCP asset URI: {file_id}")
    return f"asset://source/sha256/{digest.lower()}"


def file_id_from_asset_uri(uri: str) -> str:
    match = ASSET_URI_RE.match(uri)
    if not match:
        raise ValueError(f"unsupported asset URI: {uri}")
    return f"sha256:{match.group(1).lower()}"


def search_file_assets(
    vault_path: str | Path,
    query: str,
    *,
    db_path: Path | None = None,
    limit: int = 10,
    doc_type: str | None = None,
    source_type: str | None = None,
    product_line: str | None = None,
    customer: str | None = None,
    competitor: str | None = None,
    profile_name: str = DEFAULT_PROFILE_NAME,
    expose_paths: bool | None = None,
) -> list[dict[str, Any]]:
    vault = _vault_path(vault_path)
    registry = load_registry_by_file_id(vault)
    profile = _load_profile(vault, profile_name)
    rows = search_assets(
        vault,
        query,
        db_path=db_path,
        limit=max(limit * 3, limit),
        doc_type=doc_type,
        source_type=source_type,
        product_line=product_line,
        customer=customer,
        competitor=competitor,
        include_staged=False,
        read_only=True,
    )

    allowed: list[dict[str, Any]] = []
    for row in rows:
        record = registry.get(str(row.get("file_id") or ""))
        if record is None:
            continue
        ok, _reason = record_allowed_for_mcp(record, profile_name=profile_name, profile=profile)
        if not ok:
            continue
        allowed.append(public_asset_record({**record, **row}, expose_paths=expose_paths))
        if len(allowed) >= limit:
            break
    return allowed


def get_file_asset_metadata(
    vault_path: str | Path,
    file_id: str,
    *,
    db_path: Path | None = None,
    profile_name: str = DEFAULT_PROFILE_NAME,
    expose_paths: bool | None = None,
) -> dict[str, Any]:
    vault = _vault_path(vault_path)
    row = get_asset_by_file_id(vault, file_id, db_path=db_path, read_only=True)
    if row is None:
        raise AssetNotFoundError(f"file asset not found: {file_id}")
    registry_record = load_registry_by_file_id(vault).get(file_id)
    if registry_record is None:
        raise AssetAccessError(f"file asset is not present in source registry: {file_id}")
    ok, reason = record_allowed_for_mcp(
        registry_record,
        profile_name=profile_name,
        profile=_load_profile(vault, profile_name),
    )
    if not ok:
        raise AssetAccessError(f"file asset blocked: {reason}")
    return public_asset_record({**registry_record, **row}, expose_paths=expose_paths)


def read_file_asset_resource(
    vault_path: str | Path,
    uri: str,
    *,
    db_path: Path | None = None,
    profile_name: str = DEFAULT_PROFILE_NAME,
    max_file_bytes: int | None = None,
) -> dict[str, Any]:
    file_id = file_id_from_asset_uri(uri)
    vault = _vault_path(vault_path)
    row = get_asset_by_file_id(vault, file_id, db_path=db_path, read_only=True)
    if row is None:
        raise AssetNotFoundError(f"file asset not found: {file_id}")

    registry_record = load_registry_by_file_id(vault).get(file_id)
    if registry_record is None:
        raise AssetAccessError(f"file asset is not present in source registry: {file_id}")
    ok, reason = record_allowed_for_mcp(
        registry_record,
        profile_name=profile_name,
        profile=_load_profile(vault, profile_name),
    )
    if not ok:
        raise AssetAccessError(f"file asset blocked: {reason}")

    path = safe_asset_path(vault, str(row.get("archive_path") or ""))
    if not path.is_file():
        raise AssetNotFoundError(f"archived file missing: {row.get('archive_path')}")
    limit = DEFAULT_MAX_FILE_BYTES if max_file_bytes is None else max_file_bytes
    size = path.stat().st_size
    if limit > 0 and size > limit:
        raise AssetTooLargeError(f"file exceeds max size: {size} > {limit}")
    name = str(row.get("original_name") or path.name)
    return {
        "file_id": file_id,
        "resource_uri": uri,
        "name": name,
        "mime_type": str(row.get("mime_type") or mimetypes.guess_type(name)[0] or "application/octet-stream"),
        "size_bytes": size,
        "bytes": path.read_bytes(),
    }


def public_asset_record(record: dict[str, Any], *, expose_paths: bool | None = None) -> dict[str, Any]:
    file_id = str(record.get("file_id") or record.get("source_id") or "")
    data: dict[str, Any] = {}
    merged = {
        **record,
        "file_id": file_id,
        "resource_uri": asset_uri_for_file_id(file_id),
        "data_classification": record.get("data_classification")
        or _infer_data_classification_label(record),
    }
    for field in SAFE_RESULT_FIELDS:
        if field not in merged:
            continue
        value = _jsonish(merged[field])
        if value not in (None, "", []):
            data[field] = value
    if expose_paths if expose_paths is not None else _env_flag("OBSIDIAN_MCP_EXPOSE_PATHS", False):
        if record.get("archive_path"):
            data["archive_path"] = record["archive_path"]
        if record.get("absolute_path"):
            data["absolute_path"] = record["absolute_path"]
    return data


def record_allowed_for_mcp(
    record: dict[str, Any],
    *,
    profile_name: str = DEFAULT_PROFILE_NAME,
    profile: dict[str, Any] | None = None,
) -> tuple[bool, str]:
    doc_type = str(record.get("doc_type") or "other")
    visibility = str(record.get("visibility") or "internal").removeprefix("visibility/")
    categories = [item.lower() for item in as_list(record.get("business_categories"))]
    allowed_bots = as_list(record.get("allowed_bots"))
    excluded_bots = as_list(record.get("excluded_bots"))
    if profile_name in excluded_bots or "*" in excluded_bots:
        return False, "explicitly excluded"
    if allowed_bots and profile_name not in allowed_bots and "*" not in allowed_bots:
        return False, "not in explicit allowed_bots"
    if visibility == "pii":
        return False, "visibility pii blocked"
    if doc_type in BLOCKED_DOC_TYPES:
        return False, f"doc_type {doc_type} blocked"
    if doc_type not in ALLOWED_ENGINEERING_DOC_TYPES:
        return False, f"doc_type {doc_type} is not in engineering file whitelist"
    for category in categories:
        if any(category == prefix or category.startswith(prefix + "/") for prefix in BLOCKED_CATEGORY_PREFIXES):
            return False, f"business category {category} blocked"

    if profile is None:
        return True, "allowed"
    decision = source_decision(profile_name, _source_file_profile(profile, record), record)
    if not decision.include:
        return False, decision.reason
    return True, "allowed"


def load_registry_by_file_id(vault_path: str | Path) -> dict[str, dict[str, Any]]:
    registry_path = registry_path_for_vault(vault_path)
    if not registry_path.is_file():
        return {}
    records: dict[str, dict[str, Any]] = {}
    for raw in registry_path.read_text(encoding="utf-8", errors="ignore").splitlines():
        if not raw.strip():
            continue
        try:
            record = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if record.get("record_type") != "source_archive":
            continue
        file_id = str(record.get("source_id") or "")
        if file_id:
            records[file_id] = record
    return records


def safe_asset_path(vault_path: str | Path, archive_path: str) -> Path:
    vault = _vault_path(vault_path).resolve()
    path = Path(absolute_asset_path(vault, archive_path)).resolve()
    if path != vault and os.path.commonpath([str(vault), str(path)]) != str(vault):
        raise AssetAccessError(f"asset path escapes vault: {archive_path}")
    return path


def mcp_runtime_options(
    *,
    profile_name: str | None = None,
    max_file_bytes: int | None = None,
    expose_paths: bool | None = None,
) -> dict[str, Any]:
    return {
        "profile_name": profile_name or config_value("OBSIDIAN_MCP_BOT_PROFILE", DEFAULT_PROFILE_NAME),
        "max_file_bytes": _env_int("OBSIDIAN_MCP_MAX_FILE_BYTES", DEFAULT_MAX_FILE_BYTES)
        if max_file_bytes is None
        else max_file_bytes,
        "expose_paths": _env_flag("OBSIDIAN_MCP_EXPOSE_PATHS", False) if expose_paths is None else expose_paths,
    }


def _load_profile(vault: Path, profile_name: str) -> dict[str, Any] | None:
    try:
        return load_profiles(vault).get(profile_name)
    except FileNotFoundError:
        return None


def _source_file_profile(profile: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
    runtime_profile = {**profile, "include_source_files": True}
    data_classification = str(record.get("data_classification") or "")
    allowed = as_list(runtime_profile.get("allowed_data_classification"))
    if data_classification == "customer-shareable" and "technical-confidential" in allowed:
        runtime_profile["allowed_data_classification"] = [*allowed, "customer-shareable"]
    return runtime_profile


def _vault_path(vault_path: str | Path) -> Path:
    return Path(clean_config_value(str(vault_path))).expanduser()


def _jsonish(value: Any) -> Any:
    if isinstance(value, str):
        try:
            return json.loads(value)
        except json.JSONDecodeError:
            return value
    return value


def _infer_data_classification_label(record: dict[str, Any]) -> str:
    if str(record.get("visibility") or "").removeprefix("visibility/") == "pii":
        return "personal-sensitive"
    doc_type = str(record.get("doc_type") or "")
    if doc_type in {"drawing", "customer_3d_model", "customer_specification"}:
        return "technical-confidential"
    if doc_type == "test_report":
        return "internal"
    return "internal"


def _env_flag(key: str, default: bool) -> bool:
    raw = os.environ.get(key) or config_value(key)
    if raw == "":
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def _env_int(key: str, default: int) -> int:
    raw = os.environ.get(key) or config_value(key)
    if raw == "":
        return default
    try:
        return int(raw)
    except ValueError:
        return default