#!/usr/bin/env python3
"""
通用MCP工具包装器
为所有MCP工具提供目录访问控制和安全限制
"""

import os
import sys
import json
import asyncio
import subprocess
from typing import Dict, Any, Optional
from pathlib import Path


class MCPSecurityWrapper:
    """MCP安全包装器"""
    
    def __init__(self, allowed_directory: str):
        self.allowed_directory = os.path.abspath(allowed_directory)
        self.project_id = os.getenv("PROJECT_ID", "default")
    
    def validate_path(self, path: str) -> str:
        """验证路径是否在允许的目录内"""
        if not os.path.isabs(path):
            path = os.path.abspath(path)
        
        # 规范化路径
        path = os.path.normpath(path)
        
        # 检查路径遍历
        if ".." in path.split(os.sep):
            raise ValueError(f"路径遍历攻击被阻止: {path}")
        
        # 检查是否在允许的目录内
        if not path.startswith(self.allowed_directory):
            raise ValueError(f"访问被拒绝: {path} 不在允许的目录 {self.allowed_directory} 内")
        
        return path
    
    def safe_execute_command(self, command: list, cwd: Optional[str] = None) -> Dict[str, Any]:
        """安全执行命令，限制工作目录"""
        if cwd is None:
            cwd = self.allowed_directory
        else:
            cwd = self.validate_path(cwd)
        
        # 设置环境变量限制
        env = os.environ.copy()
        env["PWD"] = cwd
        env["PROJECT_DATA_DIR"] = self.allowed_directory
        env["PROJECT_ID"] = self.project_id
        
        try:
            result = subprocess.run(
                command,
                cwd=cwd,
                env=env,
                capture_output=True,
                text=True,
                timeout=30  # 30秒超时
            )
            
            return {
                "success": result.returncode == 0,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "returncode": result.returncode
            }
        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "stdout": "",
                "stderr": "命令执行超时",
                "returncode": -1
            }
        except Exception as e:
            return {
                "success": False,
                "stdout": "",
                "stderr": str(e),
                "returncode": -1
            }


async def handle_wrapped_request(request: Dict[str, Any], tool_name: str) -> Dict[str, Any]:
    """处理包装后的MCP请求"""
    try:
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id")
        
        allowed_dir = get_allowed_directory()
        wrapper = MCPSecurityWrapper(allowed_dir)
        
        if method == "initialize":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {}
                    },
                    "serverInfo": {
                        "name": f"{tool_name}-wrapper",
                        "version": "1.0.0"
                    }
                }
            }
        
        elif method == "ping":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "pong": True
                }
            }
        
        elif method == "tools/list":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "tools": get_tool_definitions(tool_name)
                }
            }
        
        elif method == "tools/call":
            return await execute_tool_call(wrapper, tool_name, params, request_id)
        
        else:
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": f"Unknown method: {method}"
                }
            }
    
    except Exception as e:
        return {
            "jsonrpc": "2.0",
            "id": request.get("id"),
            "error": {
                "code": -32603,
                "message": f"Internal error: {str(e)}"
            }
        }


def get_allowed_directory():
    """获取允许访问的目录"""
    project_dir = os.getenv("PROJECT_DATA_DIR", "./data")
    return os.path.abspath(project_dir)


def get_tool_definitions(tool_name: str) -> list:
    """根据工具名称返回工具定义"""
    if tool_name == "ripgrep-wrapper":
        return [
            {
                "name": "ripgrep_search",
                "description": "在项目目录内搜索文本",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "pattern": {"type": "string", "description": "搜索模式"},
                        "path": {"type": "string", "description": "搜索路径（相对于项目目录）"},
                        "maxResults": {"type": "integer", "default": 100}
                    },
                    "required": ["pattern"]
                }
            }
        ]
    elif tool_name == "directory-tree-wrapper":
        return [
            {
                "name": "get_directory_tree",
                "description": "获取项目目录结构",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "目录路径（相对于项目目录）"},
                        "max_depth": {"type": "integer", "default": 3}
                    }
                }
            }
        ]
    else:
        return []


async def execute_tool_call(wrapper: MCPSecurityWrapper, tool_name: str, params: Dict[str, Any], request_id: Any) -> Dict[str, Any]:
    """执行工具调用"""
    try:
        if tool_name == "ripgrep-wrapper":
            return await execute_ripgrep_search(wrapper, params, request_id)
        elif tool_name == "directory-tree-wrapper":
            return await execute_directory_tree(wrapper, params, request_id)
        else:
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": f"Unknown tool: {tool_name}"
                }
            }
    except Exception as e:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32603,
                "message": str(e)
            }
        }


async def execute_ripgrep_search(wrapper: MCPSecurityWrapper, params: Dict[str, Any], request_id: Any) -> Dict[str, Any]:
    """执行ripgrep搜索"""
    pattern = params.get("pattern", "")
    path = params.get("path", ".")
    max_results = params.get("maxResults", 100)
    
    # 验证和构建搜索路径
    search_path = os.path.join(wrapper.allowed_directory, path)
    search_path = wrapper.validate_path(search_path)
    
    # 构建ripgrep命令
    command = [
        "rg",
        "--json",
        "--max-count", str(max_results),
        pattern,
        search_path
    ]
    
    result = wrapper.safe_execute_command(command)
    
    if result["success"]:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "content": [
                    {
                        "type": "text",
                        "text": result["stdout"]
                    }
                ]
            }
        }
    else:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32603,
                "message": f"搜索失败: {result['stderr']}"
            }
        }


async def execute_directory_tree(wrapper: MCPSecurityWrapper, params: Dict[str, Any], request_id: Any) -> Dict[str, Any]:
    """执行目录树获取"""
    path = params.get("path", ".")
    max_depth = params.get("max_depth", 3)
    
    # 验证和构建目录路径
    dir_path = os.path.join(wrapper.allowed_directory, path)
    dir_path = wrapper.validate_path(dir_path)
    
    # 构建目录树命令
    command = [
        "find",
        dir_path,
        "-type", "d",
        "-maxdepth", str(max_depth)
    ]
    
    result = wrapper.safe_execute_command(command)
    
    if result["success"]:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "content": [
                    {
                        "type": "text",
                        "text": result["stdout"]
                    }
                ]
            }
        }
    else:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32603,
                "message": f"获取目录树失败: {result['stderr']}"
            }
        }