siyuan-mcp-server 0.1.0__tar.gz

siyuan_mcp_server-0.1.0/PKG-INFO

@@ -0,0 +1,176 @@
+Metadata-Version: 2.3
+Name: siyuan-mcp-server
+Version: 0.1.0
+Summary: 思源笔记 MCP Server - 提供思源笔记API的MCP工具接口
+Author: leolulu
+Author-email: leolulu <348699103@qq.com>
+License: MIT
+Requires-Dist: mcp
+Requires-Dist: requests
+Requires-Dist: detect-secrets
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# 思源笔记 MCP 服务器 (官方 SDK 版)
+
+本项目提供了一个基于官方 MCP Python SDK 构建的思源笔记 MCP (Model Context Protocol) 服务器。它允许 AI Agent 通过一套标准化的工具与您的思源笔记知识库进行交互。
+
+该服务器充当一座桥梁，将 MCP 的工具调用转换为对思源笔记 API 的请求，专注于提供强大的只读查询能力。
+
+## 功能特性
+
+- **基于官方 SDK 构建**: 确保了兼容性并遵循最佳实践。
+- **`FastMCP` 集成**: 使用高级的 `FastMCP` 服务器，兼具简洁与强大。
+- **生命周期管理**: 通过 `lifespan` 机制安全地管理 `SiyuanAPI` 客户端的生命周期。
+- **装饰器驱动的工具**: 使用 `@mcp.tool()` 装饰器，工具定义清晰简洁。
+- **兼具高层与底层工具**: 同时提供易于使用的高级查询工具和功能强大的底层 `execute_sql` 工具，以实现最大灵活性。
+- **敏感数据自动打码**: 自动检测并打码返回内容中的敏感信息（如 API 密钥、令牌、密码等），保护用户隐私和数据安全。
+
+## 环境要求
+
+- **Python 3.10+**（仅开发时需要，使用 uvx 运行时无需）
+- **uv**（推荐使用，用于 `uvx` 命令）
+- 思源笔记桌面客户端正在运行
+- 思源笔记 API Token（在思源笔记设置中获取）
+
+### 安装 uv
+
+如果尚未安装 uv：
+
+```bash
+# macOS/Linux
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Windows
+powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
+
+# 或使用包管理器
+pip install uv
+```
+
+## 安装与配置
+
+1.  **克隆仓库:**
+    ```bash
+    git clone <repository-url>
+    cd siyuan-mcp-server
+    ```
+
+2.  **安装依赖:**
+    我们推荐使用 `uv`。
+    ```bash
+    uv sync
+    ```
+
+
+## 如何运行
+
+### 方式一：使用 uvx（推荐，无需安装）
+
+这是最简单的方式，无需预先安装，`uvx` 会自动从 PyPI 下载并运行。
+
+**Claude Desktop 配置**:
+
+```json
+{
+  "mcpServers": {
+    "siyuan": {
+      "command": "uvx",
+      "args": ["siyuan-mcp-server"],
+      "env": {
+        "SIYUAN_API_TOKEN": "your_token_here"
+      }
+    }
+  }
+}
+```
+
+**指定版本**:
+
+```json
+{
+  "mcpServers": {
+    "siyuan": {
+      "command": "uvx",
+      "args": ["siyuan-mcp-server==0.1.0"],
+      "env": {
+        "SIYUAN_API_TOKEN": "your_token_here"
+      }
+    }
+  }
+}
+```
+
+**uvx 的优势**:
+- ✅ 无需预先安装包
+- ✅ 自动版本管理
+- ✅ 隔离的临时环境
+- ✅ 自动依赖管理
+- ✅ 快速启动（利用 uv 的缓存）
+
+### 方式二：本地开发运行
+
+在开发期间，可以使用 `uv run` 直接运行本地代码：
+
+**Claude Desktop 配置**:
+
+```json
+{
+  "mcpServers": {
+    "siyuan": {
+      "command": "uv",
+      "args": ["run", "siyuan_mcp_server"],
+      "cwd": "/path/to/siyuan-mcp-server",
+      "env": {
+        "SIYUAN_API_TOKEN": "your_token_here"
+      }
+    }
+  }
+}
+```
+
+**说明**:
+- `cwd` 指向项目根目录
+- `uv run` 会使用项目的虚拟环境
+- 代码修改后无需重新构建
+
+## 已实现的工具
+
+所有工具均在 `siyuan_mcp_server.py` 文件中定义。
+
+-   **`find_notebooks`**: 查找并列出笔记本。
+-   **`find_documents`**: 根据笔记本、标题和日期等条件查找文档。
+-   **`search_blocks`**: 根据关键词、父块、块类型和日期等条件搜索内容块。
+-   **`get_block_content`**: 获取指定块的完整 Markdown 内容。
+-   **`get_blocks_content`**: 批量获取多个块的完整内容，比多次调用 `get_block_content` 更高效。
+-   **`execute_sql`**: 直接对数据库执行只读的 `SELECT` 查询。
+
+## 未来计划
+
+- [ ] 添加更多高级查询工具
+- [ ] 支持写入操作（创建/更新文档）
+- [ ] 添加单元测试
+
+## 安全特性
+
+本项目内置了敏感数据保护机制，通过 `tools.py` 中的 `mask_sensitive_data` 函数实现：
+
+- **自动检测敏感信息**: 能够识别多种格式的敏感数据，包括：
+  - AWS Access Key ID 和 Secret Access Key
+  - GitHub Personal Access Token
+  - JWT Token
+  - UUID
+  - API Key
+  - OAuth tokens
+  - Private Key
+  - 数据库连接字符串中的密码
+  - Base64 编码的密钥
+  - 十六进制密钥
+  - 其他通用密钥格式
+
+- **智能打码策略**: 采用中间部分打码的方式，保留字符串的开头和结尾部分，便于识别但不泄露完整信息。
+
+- **全面保护**: 在所有返回用户数据的内容中自动应用打码处理，包括：
+  - 块内容搜索结果
+  - 块详细内容
+  - SQL 查询结果

siyuan_mcp_server-0.1.0/README.md

@@ -0,0 +1,163 @@
+# 思源笔记 MCP 服务器 (官方 SDK 版)
+
+本项目提供了一个基于官方 MCP Python SDK 构建的思源笔记 MCP (Model Context Protocol) 服务器。它允许 AI Agent 通过一套标准化的工具与您的思源笔记知识库进行交互。
+
+该服务器充当一座桥梁，将 MCP 的工具调用转换为对思源笔记 API 的请求，专注于提供强大的只读查询能力。
+
+## 功能特性
+
+- **基于官方 SDK 构建**: 确保了兼容性并遵循最佳实践。
+- **`FastMCP` 集成**: 使用高级的 `FastMCP` 服务器，兼具简洁与强大。
+- **生命周期管理**: 通过 `lifespan` 机制安全地管理 `SiyuanAPI` 客户端的生命周期。
+- **装饰器驱动的工具**: 使用 `@mcp.tool()` 装饰器，工具定义清晰简洁。
+- **兼具高层与底层工具**: 同时提供易于使用的高级查询工具和功能强大的底层 `execute_sql` 工具，以实现最大灵活性。
+- **敏感数据自动打码**: 自动检测并打码返回内容中的敏感信息（如 API 密钥、令牌、密码等），保护用户隐私和数据安全。
+
+## 环境要求
+
+- **Python 3.10+**（仅开发时需要，使用 uvx 运行时无需）
+- **uv**（推荐使用，用于 `uvx` 命令）
+- 思源笔记桌面客户端正在运行
+- 思源笔记 API Token（在思源笔记设置中获取）
+
+### 安装 uv
+
+如果尚未安装 uv：
+
+```bash
+# macOS/Linux
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Windows
+powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
+
+# 或使用包管理器
+pip install uv
+```
+
+## 安装与配置
+
+1.  **克隆仓库:**
+    ```bash
+    git clone <repository-url>
+    cd siyuan-mcp-server
+    ```
+
+2.  **安装依赖:**
+    我们推荐使用 `uv`。
+    ```bash
+    uv sync
+    ```
+
+
+## 如何运行
+
+### 方式一：使用 uvx（推荐，无需安装）
+
+这是最简单的方式，无需预先安装，`uvx` 会自动从 PyPI 下载并运行。
+
+**Claude Desktop 配置**:
+
+```json
+{
+  "mcpServers": {
+    "siyuan": {
+      "command": "uvx",
+      "args": ["siyuan-mcp-server"],
+      "env": {
+        "SIYUAN_API_TOKEN": "your_token_here"
+      }
+    }
+  }
+}
+```
+
+**指定版本**:
+
+```json
+{
+  "mcpServers": {
+    "siyuan": {
+      "command": "uvx",
+      "args": ["siyuan-mcp-server==0.1.0"],
+      "env": {
+        "SIYUAN_API_TOKEN": "your_token_here"
+      }
+    }
+  }
+}
+```
+
+**uvx 的优势**:
+- ✅ 无需预先安装包
+- ✅ 自动版本管理
+- ✅ 隔离的临时环境
+- ✅ 自动依赖管理
+- ✅ 快速启动（利用 uv 的缓存）
+
+### 方式二：本地开发运行
+
+在开发期间，可以使用 `uv run` 直接运行本地代码：
+
+**Claude Desktop 配置**:
+
+```json
+{
+  "mcpServers": {
+    "siyuan": {
+      "command": "uv",
+      "args": ["run", "siyuan_mcp_server"],
+      "cwd": "/path/to/siyuan-mcp-server",
+      "env": {
+        "SIYUAN_API_TOKEN": "your_token_here"
+      }
+    }
+  }
+}
+```
+
+**说明**:
+- `cwd` 指向项目根目录
+- `uv run` 会使用项目的虚拟环境
+- 代码修改后无需重新构建
+
+## 已实现的工具
+
+所有工具均在 `siyuan_mcp_server.py` 文件中定义。
+
+-   **`find_notebooks`**: 查找并列出笔记本。
+-   **`find_documents`**: 根据笔记本、标题和日期等条件查找文档。
+-   **`search_blocks`**: 根据关键词、父块、块类型和日期等条件搜索内容块。
+-   **`get_block_content`**: 获取指定块的完整 Markdown 内容。
+-   **`get_blocks_content`**: 批量获取多个块的完整内容，比多次调用 `get_block_content` 更高效。
+-   **`execute_sql`**: 直接对数据库执行只读的 `SELECT` 查询。
+
+## 未来计划
+
+- [ ] 添加更多高级查询工具
+- [ ] 支持写入操作（创建/更新文档）
+- [ ] 添加单元测试
+
+## 安全特性
+
+本项目内置了敏感数据保护机制，通过 `tools.py` 中的 `mask_sensitive_data` 函数实现：
+
+- **自动检测敏感信息**: 能够识别多种格式的敏感数据，包括：
+  - AWS Access Key ID 和 Secret Access Key
+  - GitHub Personal Access Token
+  - JWT Token
+  - UUID
+  - API Key
+  - OAuth tokens
+  - Private Key
+  - 数据库连接字符串中的密码
+  - Base64 编码的密钥
+  - 十六进制密钥
+  - 其他通用密钥格式
+
+- **智能打码策略**: 采用中间部分打码的方式，保留字符串的开头和结尾部分，便于识别但不泄露完整信息。
+
+- **全面保护**: 在所有返回用户数据的内容中自动应用打码处理，包括：
+  - 块内容搜索结果
+  - 块详细内容
+  - SQL 查询结果

siyuan_mcp_server-0.1.0/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "siyuan-mcp-server"
+version = "0.1.0"
+description = "思源笔记 MCP Server - 提供思源笔记API的MCP工具接口"
+authors = [
+    { name = "leolulu", email = "348699103@qq.com" }
+]
+license = { text = "MIT" }
+readme = "README.md"
+readme-content-type = "text/markdown"
+requires-python = ">=3.10"
+dependencies = [
+    "mcp",
+    "requests",
+    "detect-secrets",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.26,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "siyuan_mcp_server"
+
+[project.scripts]
+siyuan-mcp-server = "siyuan_mcp_server:main"

siyuan_mcp_server-0.1.0/src/siyuan_mcp_server/__init__.py

@@ -0,0 +1,324 @@
+import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mcp
+import requests
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSession
+
+from .tools import mask_sensitive_data
+
+
+# --- 1. Siyuan API Wrapper ---
+class SiyuanAPI:
+    def __init__(self, api_token: str, base_url: str = "http://127.0.0.1:6806"):
+        self.base_url = base_url
+        self.api_token = api_token
+        self.headers = {
+            "Authorization": f"Token {self.api_token}",
+            "Content-Type": "application/json",
+        }
+
+    def _post(self, endpoint: str, json_data: Optional[Dict[str, Any]] = None) -> Any:
+        url = f"{self.base_url}{endpoint}"
+        try:
+            response = requests.post(url, json=json_data, headers=self.headers)
+            response.raise_for_status()
+            api_response = response.json()
+            if api_response.get("code") != 0:
+                raise Exception(f"Siyuan API Error: {api_response.get('msg')}")
+            return api_response.get("data")
+        except requests.exceptions.RequestException as e:
+            raise ConnectionError(f"Failed to connect to Siyuan API: {e}")
+
+    def execute_sql(self, query: str) -> List[Dict[str, Any]]:
+        if not query.strip().upper().startswith("SELECT"):
+            raise ValueError("Only SELECT statements are allowed for security reasons.")
+        payload = {"stmt": query}
+        result = self._post("/api/query/sql", payload)
+        if not isinstance(result, list):
+            raise TypeError(f"Expected a list from SQL query, but got {type(result)}")
+        return result
+
+    def get_block_kramdown(self, block_id: str) -> Dict[str, Any]:
+        result = self._post("/api/block/getBlockKramdown", {"id": block_id})
+        if not isinstance(result, dict):
+            raise TypeError(f"Expected a dict for block content, but got {type(result)}")
+        return result
+
+    def list_notebooks(self) -> List[Dict[str, Any]]:
+        """获取笔记本列表"""
+        result = self._post("/api/notebook/lsNotebooks")
+        if not isinstance(result, dict) or "notebooks" not in result:
+            raise TypeError(f"Expected a dict with 'notebooks' key, but got {type(result)}")
+        return result["notebooks"]
+
+    def get_blocks_kramdown(self, block_ids: List[str]) -> List[Dict[str, Any]]:
+        """批量获取多个块的内容"""
+        results = []
+        for block_id in block_ids:
+            try:
+                result = self.get_block_kramdown(block_id)
+                results.append(result)
+            except Exception as e:
+                # 如果某个块获取失败，记录错误但继续处理其他块
+                results.append({
+                    "id": block_id,
+                    "error": str(e)
+                })
+        return results
+
+# --- 2. Application Context ---
+@dataclass
+class AppContext:
+    siyuan_api: SiyuanAPI
+
+# --- 3. Lifespan Management ---
+@asynccontextmanager
+async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
+    api_token = os.getenv("SIYUAN_API_TOKEN")
+    if not api_token:
+        raise ValueError("SIYUAN_API_TOKEN environment variable not set.")
+    
+    siyuan_api = SiyuanAPI(api_token=api_token)
+    try:
+        print("Siyuan API client initialized.")
+        yield AppContext(siyuan_api=siyuan_api)
+    finally:
+        print("Siyuan MCP Server shutting down.")
+
+# --- 4. MCP Server Instance ---
+mcp = FastMCP(
+    "siyuan-mcp-server",
+    lifespan=app_lifespan
+)
+
+# --- 5. Tool Definitions ---
+@mcp.tool()
+def find_notebooks(
+    ctx: Context[ServerSession, AppContext],
+    name: Optional[str] = None,
+    limit: int = 10
+) -> list:
+    """查找并列出思源笔记中的笔记本。
+
+    Args:
+        ctx: MCP 上下文对象，自动注入。
+        name (Optional[str]): 用于模糊搜索笔记本的名称。如果省略，则列出所有笔记本。
+        limit (int): 返回结果的最大数量，默认为 10。
+
+    Returns:
+        list: 包含笔记本信息的字典列表，每个字典包含 'name' 和 'id'。
+    """
+    api = ctx.request_context.lifespan_context.siyuan_api
+    notebooks = api.list_notebooks()
+    
+    # 如果指定了名称，则进行过滤
+    if name:
+        notebooks = [nb for nb in notebooks if name.lower() in nb.get("name", "").lower()]
+    
+    # 限制返回结果数量
+    return notebooks[:limit]
+
+@mcp.tool()
+def find_documents(
+    ctx: Context[ServerSession, AppContext],
+    notebook_id: Optional[str] = None,
+    title: Optional[str] = None,
+    created_after: Optional[str] = None,
+    updated_after: Optional[str] = None,
+    limit: int = 10,
+) -> list:
+    """在指定的笔记本中查找文档，支持多种过滤条件。
+
+    Args:
+        ctx: MCP 上下文对象，自动注入。
+        notebook_id (Optional[str]): 在哪个笔记本中查找。如果省略，则在所有打开的笔记本中查找。
+        title (Optional[str]): 根据文档标题进行模糊匹配。
+        created_after (Optional[str]): 查找在此日期之后创建的文档，格式为 'YYYYMMDDHHMMSS'。
+        updated_after (Optional[str]): 查找在此日期之后更新的文档，格式为 'YYYYMMDDHHMMSS'。
+        limit (int): 返回结果的最大数量，默认为 10。
+
+    Returns:
+        list: 包含文档信息的字典列表，每个字典包含 'name', 'id', 和 'hpath'。
+    """
+    api = ctx.request_context.lifespan_context.siyuan_api
+    query = "SELECT name, id, hpath FROM blocks WHERE type = 'd'"
+    conditions = []
+    if notebook_id:
+        sanitized_id = notebook_id.replace("'", "''")
+        conditions.append(f"box = '{sanitized_id}'")
+    if title:
+        sanitized_title = title.replace("'", "''")
+        conditions.append(f"name LIKE '%{sanitized_title}%'")
+    if created_after:
+        sanitized_date = created_after.replace("'", "''")
+        conditions.append(f"created > '{sanitized_date}'")
+    if updated_after:
+        sanitized_date = updated_after.replace("'", "''")
+        conditions.append(f"updated > '{sanitized_date}'")
+    if conditions:
+        query += " AND " + " AND ".join(conditions)
+    query += f" LIMIT {limit}"
+    return api.execute_sql(query)
+
+@mcp.tool()
+def search_blocks(
+    ctx: Context[ServerSession, AppContext],
+    query: str,
+    parent_id: Optional[str] = None,
+    block_type: Optional[str] = None,
+    created_after: Optional[str] = None,
+    updated_after: Optional[str] = None,
+    limit: int = 20,
+) -> list:
+    """根据关键词、类型等多种条件在思源笔记中搜索内容块。
+
+    这是最核心和最灵活的查询工具。
+
+    Args:
+        ctx: MCP 上下文对象，自动注入。
+        query (str): 在块内容中搜索的关键词。
+        parent_id (Optional[str]): 在哪个文档或父块下进行搜索。如果省略，则全局搜索。
+        block_type (Optional[str]): 限制块的类型，例如 'p' (段落), 'h' (标题), 'l' (列表)。
+        created_after (Optional[str]): 查找在此日期之后创建的块，格式为 'YYYYMMDDHHMMSS'。
+        updated_after (Optional[str]): 查找在此日期之后更新的块，格式为 'YYYYMMDDHHMMSS'。
+        limit (int): 返回结果的最大数量，默认为 20。
+
+    Returns:
+        list: 包含块信息的字典列表。
+    """
+    api = ctx.request_context.lifespan_context.siyuan_api
+    sql_query = "SELECT id, content, type, subtype, hpath FROM blocks WHERE content LIKE ?"
+    params = [f"%{query}%"]
+    if parent_id:
+        sql_query += " AND parent_id = ?"
+        params.append(parent_id)
+    if block_type:
+        sql_query += " AND type = ?"
+        params.append(block_type)
+    if created_after:
+        sql_query += " AND created > ?"
+        params.append(created_after)
+    if updated_after:
+        sql_query += " AND updated > ?"
+        params.append(updated_after)
+    sql_query += f" LIMIT {limit}"
+    for param in params:
+        sanitized_param = str(param).replace("'", "''")
+        sql_query = sql_query.replace("?", f"'{sanitized_param}'", 1)
+    results = api.execute_sql(sql_query)
+    
+    # 对搜索结果中的内容进行打码处理
+    for result in results:
+        if isinstance(result, dict):
+            if "content" in result:
+                result["content"] = mask_sensitive_data(result["content"])
+    
+    return results
+
+@mcp.tool()
+def get_block_content(
+    ctx: Context[ServerSession, AppContext],
+    block_id: str
+) -> dict:
+    """获取指定 ID 的块的完整内容。
+
+    在通过 search_blocks 找到相关块后，使用此工具读取其详细内容。
+
+    Args:
+        ctx: MCP 上下文对象，自动注入。
+        block_id (str): 要获取内容的块的 ID。
+
+    Returns:
+        dict: 包含块 Kramdown 源码等信息的字典。
+    """
+    api = ctx.request_context.lifespan_context.siyuan_api
+    result = api.get_block_kramdown(block_id)
+    
+    # 对内容进行打码处理
+    if isinstance(result, dict) and "kramdown" in result:
+        result["kramdown"] = mask_sensitive_data(result["kramdown"])
+    if isinstance(result, dict) and "content" in result:
+        result["content"] = mask_sensitive_data(result["content"])
+    
+    return result
+
+@mcp.tool()
+def get_blocks_content(
+    ctx: Context[ServerSession, AppContext],
+    block_ids: List[str]
+) -> list:
+    """批量获取多个块的完整内容。
+
+    在通过 find_documents 或 search_blocks 找到相关块后，使用此工具批量读取它们的详细内容。
+    相比多次调用 get_block_content，这个工具更高效，特别适合查询大量块。
+
+    Args:
+        ctx: MCP 上下文对象，自动注入。
+        block_ids (List[str]): 要获取内容的块的 ID 列表。
+
+    Returns:
+        list: 包含多个块信息的字典列表，每个字典包含块的 Kramdown 源码等信息。
+    """
+    api = ctx.request_context.lifespan_context.siyuan_api
+    results = api.get_blocks_kramdown(block_ids)
+    
+    # 对每个块的内容进行打码处理
+    for result in results:
+        if isinstance(result, dict):
+            if "kramdown" in result:
+                result["kramdown"] = mask_sensitive_data(result["kramdown"])
+            if "content" in result:
+                result["content"] = mask_sensitive_data(result["content"])
+    
+    return results
+
+@mcp.tool()
+def execute_sql(
+    ctx: Context[ServerSession, AppContext],
+    query: str
+) -> list:
+    """直接执行一条只读的 SQL 查询语句。
+
+    这是一个强大的底层工具，仅用于高级或复杂的查询场景。
+    为了安全，此工具只允许执行 'SELECT' 语句。
+
+    Args:
+        ctx: MCP 上下文对象，自动注入。
+        query (str): 要执行的 SQL 'SELECT' 语句。
+
+    Returns:
+        list: 代表查询结果的字典列表。
+    """
+    api = ctx.request_context.lifespan_context.siyuan_api
+    results = api.execute_sql(query)
+    
+    # 对查询结果中的内容进行打码处理
+    for result in results:
+        if isinstance(result, dict):
+            # 对可能包含敏感信息的字段进行打码处理
+            for key, value in result.items():
+                if isinstance(value, str) and len(value) > 10:
+                    # 对长字符串进行打码处理，避免误判
+                    result[key] = mask_sensitive_data(value)
+    
+    return results
+
+# --- 6. Server Runner ---
+def main():
+    """MCP 服务器入口函数
+
+    通过 uvx 或 pip install 安装后，可通过命令行直接运行此服务器。
+    """
+    mcp.run()
+
+if __name__ == "__main__":
+    # 运行方式:
+    # 1. 设置 SIYUAN_API_TOKEN 环境变量
+    # 2. 使用 uv run 直接运行: uv run siyuan_mcp_server.py
+    # 3. 安装后使用: siyuan-mcp-server
+    main()

siyuan_mcp_server-0.1.0/src/siyuan_mcp_server/tools.py

@@ -0,0 +1,92 @@
+import re
+
+
+def mask_middle_third(text):
+    """
+    只打码字符串中间1/3的部分，保留开头和结尾部分
+
+    参数:
+        text (str): 输入的字符串
+
+    返回:
+        str: 处理后的字符串，中间1/3部分被替换为*
+    """
+    if len(text) < 6:  # 如果字符串太短，直接全部打码
+        return "*" * len(text)
+
+    # 计算各个部分的长度
+    third = len(text) // 3
+    start_length = (len(text) - third) // 2
+    end_length = len(text) - third - start_length
+
+    # 构建结果字符串
+    result = text[:start_length] + ("*" * third) + text[-end_length:] if end_length > 0 else text[:start_length] + ("*" * third)
+
+    return result
+
+
+def mask_sensitive_data(text):
+    """
+    对文本中的敏感信息（密钥、API Key、Secret等）进行打码处理
+
+    参数:
+        text (str): 输入的文本
+
+    返回:
+        str: 处理后的文本，其中敏感信息被替换为*
+    """
+    # 定义各种密钥格式的正则表达式模式
+    patterns = [
+        # AWS Access Key ID: AKIA开头，20个字符
+        (r"AKIA[0-9A-Z]{16}", lambda m: mask_middle_third(m.group())),
+        # AWS Secret Access Key: 40个字符的随机字符串
+        (r"[A-Za-z0-9/+=]{40}", lambda m: mask_middle_third(m.group())),
+        # GitHub Personal Access Token
+        (
+            r"ghp_[a-zA-Z0-9]{36}|gho_[a-zA-Z0-9]{36}|ghu_[a-zA-Z0-9]{36}|ghs_[a-zA-Z0-9]{36}|ghr_[a-zA-Z0-9]{36}",
+            lambda m: mask_middle_third(m.group()),
+        ),
+        # JWT Token: 由三部分组成，用点分隔
+        (r"[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+", lambda m: mask_middle_third(m.group())),
+        # UUID
+        (r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}", lambda m: mask_middle_third(m.group())),
+        # API Key: 32位以上的字母数字组合
+        (r"[A-Za-z0-9]{32,}", lambda m: mask_middle_third(m.group())),
+        # OAuth tokens: 20位以上的字母数字组合
+        (r"[A-Za-z0-9]{20,}", lambda m: mask_middle_third(m.group())),
+        # Private Key
+        (r"-----BEGIN(?: RSA)? PRIVATE KEY-----.*?-----END(?: RSA)? PRIVATE KEY-----", lambda m: mask_middle_third(m.group())),
+        # Database URLs - 特殊处理，只打码密码部分
+        (
+            r"(postgresql|mysql|mongodb)://([^:]+):([^@]+)@([^/]+)/([^\s]+)",
+            lambda m: f"{m.group(1)}://{m.group(2)}:{mask_middle_third(m.group(3))}@{m.group(4)}/{m.group(5)}",
+        ),
+        # API URLs with credentials - 特殊处理，只打码密钥值部分
+        (r"(api[_-]?key[=:\s]+)([^\s&]+)", lambda m: f"{m.group(1)}{mask_middle_third(m.group(2))}"),
+        # Base64编码的密钥
+        (r"[A-Za-z0-9+/]{20,}={0,2}", lambda m: mask_middle_third(m.group())),
+        # 十六进制密钥
+        (r"[0-9a-fA-F]{32,}", lambda m: mask_middle_third(m.group())),
+        # 带有引号的密钥
+        (r"([\"\'])([A-Za-z0-9+/=]{20,})(\1)", lambda m: m.group(1) + mask_middle_third(m.group(2)) + m.group(3)),
+        # 通用密钥格式：包含特殊字符的长字符串
+        (r"[\"\']?[A-Za-z0-9_\-+/=]{20,}[\"\']?", lambda m: mask_middle_third(m.group())),
+    ]
+
+    # 应用所有模式
+    result = text
+    for pattern, replacement in patterns:
+        result = re.sub(pattern, replacement, result, flags=re.DOTALL)
+
+    return result
+
+
+if __name__ == "__main__":
+    test_string = """
+- Todoist token
+数据库连接: jdbc:mysql://localhost:3306/mydb?user=admin&password=secret123
+API密钥: sk_test_1234567890abcdefghijklmnopqrstuvwxyz
+令牌: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c
+用户名: user1, 密码: P@ssw0rd!
+"""
+    print(mask_sensitive_data(test_string))