mattermost-dl/mm_search_mcp.py

#!/usr/bin/env python3
"""MCP server (stdio) exposing mm_search.py search capabilities.

Cursor / Claude Desktop config example (~/.cursor/mcp.json):

    {
      "mcpServers": {
        "mattermost-search": {
          "command": "/path/to/mattermost-dl/.venv/bin/python",
          "args": ["/path/to/mattermost-dl/mm_search_mcp.py"],
          "cwd": "/path/to/mattermost-dl",
          "env": {
            "MM_SEARCH_ROOT": "results",
            "MM_SEARCH_DB": "mattermost_index.sqlite"
          }
        }
      }
    }

Requires: pip install fastmcp
"""
from __future__ import annotations

import datetime as dt
import os
from typing import Any, Literal, Optional

from fastmcp import FastMCP

import mm_search as mm_search_lib
import mm_users

mcp = FastMCP("mattermost-search")

SearchMode = Literal["message", "file_name", "file_content"]


def _env_root() -> str:
    return os.environ.get("MM_SEARCH_ROOT", "results")


def _env_db() -> str:
    return os.environ.get("MM_SEARCH_DB", "mattermost_index.sqlite")


def _parse_date(value: str) -> dt.datetime:
    try:
        return mm_search_lib.parse_iso_date(value)
    except Exception as exc:
        raise ValueError(f"invalid date '{value}': {exc}") from exc


def _collect_matches(
    *,
    query: str,
    root: str,
    user: str,
    channel: str,
    team: str,
    after: dt.datetime,
    before: dt.datetime,
    regex: bool,
    search_mode: SearchMode,
    ignore_case: bool,
    limit: int = 10,
    db: str,
    no_index: bool,
) -> list[dict[str, Any]]:
    if search_mode != "message" and regex:
        raise ValueError("search_mode file_name/file_content do not support regex=True")

    effective_ignore_case = ignore_case if ignore_case is not None else (not regex)
    matcher = mm_search_lib.make_matcher(query, regex, effective_ignore_case)
    use_index = (not no_index) and os.path.exists(db)

    if search_mode in ("file_name", "file_content"):
        if not use_index:
            raise ValueError(
                "Attachment search (file_name/file_content) requires sqlite index. "
                "Create it with mm_index.py and do not set no_index=True."
            )
        mode = "name" if search_mode == "file_name" else "content"
        matches_iter = mm_search_lib.search_attachments_index(
            db_path=db,
            query=query,
            user=user,
            channel=channel,
            team=team,
            after=after,
            before=before,
            limit=limit,
            mode=mode,
        )
        output_limit = None
    elif use_index and not regex:
        matches_iter = mm_search_lib.search_index(
            db_path=db,
            query=query,
            user=user,
            channel=channel,
            team=team,
            after=after,
            before=before,
            limit=limit,
            ignore_case=effective_ignore_case,
        )
        output_limit = None
    else:
        files = list(mm_search_lib.iter_result_files(root))
        matches_iter = (
            m
            for path in files
            for m in mm_search_lib.search_file(
                path=path,
                matcher=matcher,
                user=user,
                channel=channel,
                team=team,
                after=after,
                before=before,
            )
        )
        output_limit = limit

    out: list[dict[str, Any]] = []
    count = 0
    for match in matches_iter:
        if output_limit is not None and count >= output_limit:
            break
        out.append(match)
        count += 1
    return out


def _format_human(matches: list[dict[str, Any]], show_file: bool) -> str:
    lines: list[str] = []
    for m in matches:
        p = m["post"]
        created = p.get("created", "")
        username = p.get("username", "")
        msg = (p.get("message") or "").replace("\n", "\\n")
        attachment = m.get("attachment")
        attachment_suffix = ""
        if isinstance(attachment, dict) and attachment.get("filename"):
            attachment_suffix = f" [file: {attachment['filename']}]"
        if show_file:
            file_rel = os.path.relpath(m["file"])
            idx = p.get("idx", "?")
            prefix = f"{file_rel}:{idx}: "
        else:
            prefix = ""
        channel = m.get("channel") or "?"
        team = m.get("team") or ""
        team_prefix = f"[{team}]" if team else ""
        lines.append(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
    return "\n".join(lines)


@mcp.tool
def mm_search(
    query: str,
    root: str = None,
    user: str = None,
    channel: str = None,
    team: str = None,
    after: str = None,
    before: str = None,
    regex: bool = False,
    search_mode: SearchMode = "message",
    ignore_case: bool = None,
    limit: int = 10,
    db: str = None,
    no_index: bool = False,
    output: Literal["json", "human"] = "json",
    show_file: bool = False,
) -> dict[str, Any]:
    """Search Mattermost export JSON files produced by mattermost-dl.py.

    Mirrors mm_search.py CLI. Returns structured matches or human-readable lines.

    Args:
        query: Search string or regex (when regex=True).
        root: Root directory with export results (default: MM_SEARCH_ROOT or "results").
        user: Filter by username (exact match).
        channel: Filter by channel display_name or name (exact match).
        team: Filter by team name (exact match).
        after: Only messages on/after this date (YYYY-MM-DD or ISO).
        before: Only messages on/before this date (YYYY-MM-DD or ISO).
        regex: Treat query as regular expression (message scan only).
        search_mode: "message" (default), "file_name", or "file_content".
        ignore_case: Case-insensitive search. Default: True for substring, False for regex.
        limit: Maximum number of matches to return (default: 10).
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
        no_index: Do not use sqlite index even if DB exists.
        output: "json" returns match objects; "human" returns CLI-style text lines.
        show_file: Include source filename and per-file index in human output.
    """
    resolved_root = root or _env_root()
    resolved_db = db or _env_db()
    parsed_after = _parse_date(after) if after else None
    parsed_before = _parse_date(before) if before else None

    matches = _collect_matches(
        query=query,
        root=resolved_root,
        user=user,
        channel=channel,
        team=team,
        after=parsed_after,
        before=parsed_before,
        regex=regex,
        search_mode=search_mode,
        ignore_case=ignore_case,
        limit=limit,
        db=resolved_db,
        no_index=no_index,
    )

    result: dict[str, Any] = {
        "count": len(matches),
        "used_index": (not no_index)
        and os.path.exists(resolved_db)
        and (search_mode in ("file_name", "file_content") or not regex),
        "search_mode": search_mode,
        "root": resolved_root,
        "db": resolved_db,
    }

    if output == "human":
        result["text"] = _format_human(matches, show_file=show_file)
    else:
        result["matches"] = matches

    return result


@mcp.tool
def mm_resolve_user(
    ref: str,
    db: str = None,
) -> dict[str, Any]:
    """Resolve display name, FIO, nickname, or username to Mattermost username.

    Args:
        ref: User reference — username, first/last name, nickname, or FIO
             (e.g. "Цисык", "vladislav.tsisyk").
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
    """
    resolved_db = db or _env_db()
    return mm_users.resolve_user(resolved_db, ref)


@mcp.tool
def mm_search_users(
    query: str,
    limit: int = 10,
    db: str = None,
) -> dict[str, Any]:
    """Search users in the local directory by name or username.

    Args:
        query: Substring to match against username, first/last name, nickname, full name.
        limit: Maximum matches (default: 10).
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
    """
    resolved_db = db or _env_db()
    matches = mm_users.search_users(resolved_db, query, limit=limit)
    return {
        "query": query,
        "count": len(matches),
        "db": resolved_db,
        "matches": matches,
    }


@mcp.tool
def mm_user_directory_refresh(
    root: str = None,
    db: str = None,
    from_api: bool = False,
    config_path: str = None,
) -> dict[str, Any]:
    """Rebuild the local user directory from users.json, Mattermost API, and post usernames.

    Args:
        root: Export root (default: MM_SEARCH_ROOT or "results").
        db: SQLite index path (default: MM_SEARCH_DB).
        from_api: Also fetch profiles from Mattermost API (needs config.json with token).
        config_path: Path to Mattermost config (default: MM_CONFIG_PATH or "config.json").
    """
    resolved_root = root or _env_root()
    resolved_db = db or _env_db()
    return mm_users.refresh_user_directory(
        resolved_db,
        resolved_root,
        config_path=config_path,
        from_api=from_api,
    )


@mcp.tool
def mm_search_index_status(
    db: Optional[str] = None,
    root: Optional[str] = None,
) -> dict[str, Any]:
    """Report whether the sqlite index and export root are available."""
    resolved_db = db or _env_db()
    resolved_root = root or _env_root()
    db_exists = os.path.exists(resolved_db)
    root_exists = os.path.isdir(resolved_root)

    json_count = 0
    if root_exists:
        json_count = sum(1 for _ in mm_search_lib.iter_result_files(resolved_root))

    user_count = 0
    if db_exists:
        try:
            import sqlite3

            conn = sqlite3.connect(resolved_db)
            mm_users.ensure_users_schema(conn)
            user_count = conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]
            conn.close()
        except sqlite3.Error:
            user_count = None

    return {
        "db": resolved_db,
        "db_exists": db_exists,
        "root": resolved_root,
        "root_exists": root_exists,
        "json_files": json_count,
        "users": user_count,
        "index_required_for": ["file_name", "file_content"],
    }


if __name__ == "__main__":
    mcp.run(transport="stdio")