#!/usr/bin/env python3
"""
Code Interpreter MCP Server
Use/Using标准MCP协议实现的安全Python代码执行环境
"""

import os
import sys
import io
import tempfile
import subprocess
import contextlib
import json
import asyncio
from typing import Any, Dict, List, Optional, Sequence
from pathlib import Path
import ast
import traceback
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

# Create MCP server实例
server = Server("code-interpreter-mcp")

# 安全的内置函数白名单
SAFE_BUILTINS = {
    'abs', 'all', 'any', 'ascii', 'bin', 'bool', 'bytearray', 'bytes',
    'callable', 'chr', 'classmethod', 'complex', 'dict', 'dir', 'divmod',
    'enumerate', 'eval', 'filter', 'float', 'format', 'frozenset',
    'getattr', 'globals', 'hasattr', 'hash', 'hex', 'id', 'int',
    'isinstance', 'issubclass', 'iter', 'len', 'list', 'locals',
    'map', 'max', 'min', 'next', 'object', 'oct', 'ord', 'pow',
    'property', 'range', 'repr', 'reversed', 'round', 'set', 'setattr',
    'slice', 'sorted', 'staticmethod', 'str', 'sum', 'super', 'tuple',
    'type', 'vars', 'zip', 'print', '__import__', 'open', 'compile',
    'exec', 'eval', 'input', 'help', 'quit', 'exit', 'copyright',
    'credits', 'license', 'breakpoint', 'memoryview', 'delattr'
}

# 允许的模块白名单 - 大幅扩展
ALLOWED_MODULES = {
    # 基础模块
    'math', 'random', 'datetime', 'json', 'base64', 'urllib', 'urllib.parse',
    'collections', 'itertools', 'functools', 'operator', 'string',
    'decimal', 'fractions', 'statistics', 'uuid', 'hashlib', 'hmac',
    'secrets', 'time', 'calendar', 'zoneinfo', 'dataclasses', 'typing',
    'enum', 'abc', 'copy', 'pickle', 'shelve', 'sqlite3', 'csv',
    're', 'regex', 'difflib', 'textwrap', 'unicodedata', 'stringprep',
    'struct', 'codecs', 'binascii', 'array', 'weakref', 'types',
    'pprint', 'reprlib', 'contextvars', 'contextlib', 'atexit',
    'traceback', 'linecache', 'timeit', 'trace', 'dis', 'inspect',
    
    # File和系统
    'os', 'sys', 'io', 'pathlib', 'tempfile', 'glob', 'fnmatch',
    'shutil', 'zipfile', 'tarfile', 'gzip', 'bz2', 'lzma', 'zlib',
    'configparser', 'argparse', 'getopt', 'logging', 'platform',
    'errno', 'ctypes', 'subprocess', 'multiprocessing', 'concurrent',
    'concurrent.futures', 'threading', 'queue', 'asyncio', 'socket',
    'ssl', 'select', 'selectors', 'signal', 'mmap', 'resource',
    
    # 网络和互联网
    'urllib.request', 'urllib.error', 'urllib.robotparser', 'http',
    'http.client', 'ftplib', 'poplib', 'imaplib', 'smtplib', 'email',
    'mailbox', 'mimetypes', 'webbrowser', 'wsgiref', 'xmlrpc',
    'ipaddress', 'socketserver', 'http.server', 'http.cookies',
    'http.cookiejar', 'xml', 'xml.etree', 'xml.etree.ElementTree',
    'xml.dom', 'xml.sax', 'html', 'html.parser', 'html.entities',
    
    # Data科学和机器学习
    'numpy', 'pandas', 'matplotlib', 'matplotlib.pyplot', 'matplotlib.patches',
    'seaborn', 'scipy', 'scipy.stats', 'scipy.optimize', 'scipy.signal',
    'sklearn', 'sklearn.linear_model', 'sklearn.model_selection', 'sklearn.metrics',
    'sklearn.preprocessing', 'sklearn.ensemble', 'sklearn.tree', 'sklearn.cluster',
    'sklearn.decomposition', 'sklearn.feature_extraction', 'sklearn.pipeline',
    'statsmodels', 'xgboost', 'lightgbm', 'catboost', 'tensorflow', 'keras',
    'torch', 'torchvision', 'transformers', 'cv2', 'opencv-python',
    
    # DataHandle/Process和分析
    'requests', 'beautifulsoup4', 'bs4', 'lxml', 'scrapy', 'selenium',
    'openpyxl', 'xlrd', 'xlwt', 'xlsxwriter', 'pyarrow', 'fastparquet',
    'h5py', 'netCDF4', 'xarray', 'dask', 'joblib', 'numba', 'cython',
    
    # 可视化
    'plotly', 'plotly.graph_objects', 'plotly.express', 'bokeh', 'altair',
    'holoviews', 'dash', 'streamlit', 'gradio', 'panel', 'hvplot',
    'folium', 'geopandas', 'shapely', 'pyproj', 'cartopy', 'basemap',
    
    # 文本Handle/Process和NLP
    'nltk', 'spacy', 'textblob', 'gensim', 'wordcloud', 'fuzzywuzzy',
    'python-Levenshtein', 'langdetect', 'jieba', 'pypinyin', 'opencc',
    
    # 图像Handle/Process
    'PIL', 'PIL.Image', 'pillow', 'imageio', 'scikit-image', 'skimage',
    'wand', 'pytesseract', 'qrcode', 'barcode', 'cairosvg', 'svgwrite',
    
    # PDF和文档Handle/Process
    'reportlab', 'reportlab.lib', 'reportlab.lib.pagesizes', 'reportlab.lib.styles',
    'reportlab.lib.units', 'reportlab.platypus', 'reportlab.pdfgen', 'reportlab.pdfgen.canvas',
    'PyPDF2', 'pdfplumber', 'fpdf', 'fpdf2', 'python-docx', 'docx',
    'python-pptx', 'pptx', 'markdown', 'textile', 'docutils', 'sphinx',
    
    # Data库
    'pymongo', 'redis', 'psycopg2', 'mysql-connector-python', 'pymysql',
    'sqlalchemy', 'peewee', 'dataset', 'tinydb', 'pickledb',
    
    # 其他实用工具
    'click', 'typer', 'rich', 'tqdm', 'colorama', 'termcolor', 'tabulate',
    'prettytable', 'humanize', 'python-dateutil', 'pytz', 'arrow',
    'pendulum', 'schedule', 'apscheduler', 'celery', 'rq', 'huey',
    'pydantic', 'marshmallow', 'attrs', 'cattrs', 'cerberus', 'voluptuous',
    'jsonschema', 'pyyaml', 'toml', 'configobj', 'python-dotenv',
    'watchdog', 'psutil', 'py-cpuinfo', 'netifaces', 'pythonping',
    'paramiko', 'fabric', 'invoke', 'sh', 'delegator.py', 'envoy',
    'pytest', 'unittest', 'nose', 'mock', 'faker', 'factory_boy',
    'hypothesis', 'coverage', 'pytest-cov', 'tox', 'nox', 'pre-commit'
}

def is_safe_code(code: str) -> tuple[bool, str]:
    """检查代码是否安全 - 放宽Limit，只阻止最危险的操作"""
    try:
        tree = ast.parse(code)
        
        # 只检查最危险的操作
        for node in ast.walk(tree):
            # 允许所有导入（移除模块白名单检查）
            # 用户应该能够导入他们Need/Require的任何模块
            
            # 只禁止访问最危险的底层属性
            if isinstance(node, ast.Attribute):
                dangerous_attrs = ['__builtins__', '__globals__', '__locals__']
                if node.attr in dangerous_attrs:
                    return False, f"不允许访问底层属性: {node.attr}"
            
            # 不再禁止任何函数调用
            # 允许 open, exec, eval, compile 等所有函数
        
        return True, "代码安全"
    except SyntaxError as e:
        return False, f"语法Error: {e}"

@server.list_tools()
async def list_tools() -> List[Tool]:
    """List all available tools"""
    return [
        Tool(
            name="execute_python_code",
            description="执行Python代码",
            inputSchema={
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "要执行的Python代码"},
                    "timeout": {"type": "integer", "description": "超时时间（秒）", "default": 30}
                },
                "required": ["code"]
            }
        ),
        Tool(
            name="create_plot",
            description="创建图表并保存",
            inputSchema={
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "绘图代码"},
                    "output_format": {"type": "string", "description": "输出格式 (png, jpg, svg, pdf)", "default": "png"}
                },
                "required": ["code"]
            }
        ),
        Tool(
            name="analyze_data",
            description="Data分析工具",
            inputSchema={
                "type": "object",
                "properties": {
                    "data_code": {"type": "string", "description": "Data准备代码"},
                    "analysis_code": {"type": "string", "description": "分析代码"}
                },
                "required": ["data_code", "analysis_code"]
            }
        ),
        Tool(
            name="install_package",
            description="安装Python包 (仅限白名单)",
            inputSchema={
                "type": "object",
                "properties": {
                    "package_name": {"type": "string", "description": "包名"}
                },
                "required": ["package_name"]
            }
        )
    ]

@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> Sequence[TextContent]:
    """Call tool"""
    try:
        if name == "execute_python_code":
            result = await execute_python_code_impl(
                arguments["code"], 
                arguments.get("timeout", 30)
            )
        elif name == "create_plot":
            result = await create_plot_impl(
                arguments["code"], 
                arguments.get("output_format", "png")
            )
        elif name == "analyze_data":
            result = await analyze_data_impl(
                arguments["data_code"], 
                arguments["analysis_code"]
            )
        elif name == "install_package":
            result = await install_package_impl(arguments["package_name"])
        else:
            result = {"status": "error", "message": f"未知工具: {name}"}
        
        return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))]
    
    except Exception as e:
        error_result = {"status": "error", "message": str(e)}
        return [TextContent(type="text", text=json.dumps(error_result, indent=2, ensure_ascii=False))]

async def execute_python_code_impl(code: str, timeout: int = 30) -> Dict[str, Any]:
    """执行Python代码"""
    try:
        # 安全性检查
        is_safe, message = is_safe_code(code)
        if not is_safe:
            return {
                "status": "error",
                "message": f"代码安全检查Failed: {message}",
                "code": code
            }
        
        # 创建执行环境 - Use/Using完整的内置函数
        safe_globals = {
            '__builtins__': __builtins__  # Use/Using完整的内置函数，不再Limit
        }
        
        # 添加允许的模块
        import math
        import random
        import datetime
        import json
        import base64
        import urllib.parse
        import collections
        import itertools
        import functools
        import operator
        import string
        import decimal
        import fractions
        import statistics
        import uuid
        import hashlib
        
        safe_globals.update({
            'math': math,
            'random': random,
            'datetime': datetime,
            'json': json,
            'base64': base64,
            'urllib': urllib,
            'collections': collections,
            'itertools': itertools,
            'functools': functools,
            'operator': operator,
            'string': string,
            'decimal': decimal,
            'fractions': fractions,
            'statistics': statistics,
            'uuid': uuid,
            'hashlib': hashlib
        })
        
        # 尝试添加Data科学库
        try:
            import numpy as np
            safe_globals['numpy'] = np
            safe_globals['np'] = np
        except ImportError:
            pass
            
        try:
            import pandas as pd
            safe_globals['pandas'] = pd
            safe_globals['pd'] = pd
        except ImportError:
            pass
            
        try:
            import matplotlib.pyplot as plt
            import matplotlib
            safe_globals['matplotlib'] = matplotlib
            safe_globals['plt'] = plt
        except ImportError:
            pass
            
        try:
            import seaborn as sns
            safe_globals['seaborn'] = sns
            safe_globals['sns'] = sns
        except ImportError:
            pass
            
        try:
            import io
            safe_globals['io'] = io
        except ImportError:
            pass
            
        try:
            from io import BytesIO
            safe_globals['BytesIO'] = BytesIO
        except ImportError:
            pass
        
        # 捕获输出
        output_buffer = io.StringIO()
        
        with contextlib.redirect_stdout(output_buffer):
            with contextlib.redirect_stderr(output_buffer):
                # 执行代码
                exec(code, safe_globals)
        
        output = output_buffer.getvalue()
        
        return {
            "status": "success",
            "output": output,
            "code": code,
            "message": "代码执行Success"
        }
    
    except Exception as e:
        error_traceback = traceback.format_exc()
        return {
            "status": "error",
            "message": str(e),
            "traceback": error_traceback,
            "code": code
        }

async def create_plot_impl(code: str, output_format: str = "png") -> Dict[str, Any]:
    """创建图表并保存"""
    try:
        # 安全性检查
        is_safe, message = is_safe_code(code)
        if not is_safe:
            return {
                "status": "error",
                "message": f"代码安全检查Failed: {message}",
                "code": code
            }
        
        # 创建临时File
        with tempfile.NamedTemporaryFile(suffix=f'.{output_format}', delete=False) as tmp_file:
            temp_path = tmp_file.name
        
        # 准备绘图环境
        safe_globals = {
            '__builtins__': {name: getattr(__builtins__, name) for name in SAFE_BUILTINS if hasattr(__builtins__, name)}
        }
        
        # 导入绘图库
        try:
            import matplotlib.pyplot as plt
            import numpy as np
            
            safe_globals.update({
                'plt': plt,
                'numpy': np,
                'np': np
            })
            
            # 尝试导入其他常用库
            try:
                import pandas as pd
                safe_globals['pandas'] = pd
                safe_globals['pd'] = pd
            except ImportError:
                pass
            
            try:
                import seaborn as sns
                safe_globals['seaborn'] = sns
                safe_globals['sns'] = sns
            except ImportError:
                pass
            
        except ImportError:
            return {
                "status": "error",
                "message": "matplotlib未安装，无法创建图表"
            }
        
        # 执行绘图代码
        output_buffer = io.StringIO()
        
        with contextlib.redirect_stdout(output_buffer):
            with contextlib.redirect_stderr(output_buffer):
                exec(code, safe_globals)
                plt.savefig(temp_path, format=output_format, dpi=300, bbox_inches='tight')
                plt.close()
        
        output = output_buffer.getvalue()
        
        return {
            "status": "success",
            "output": output,
            "plot_path": temp_path,
            "format": output_format,
            "code": code,
            "message": "图表创建Success"
        }
    
    except Exception as e:
        error_traceback = traceback.format_exc()
        return {
            "status": "error",
            "message": str(e),
            "traceback": error_traceback,
            "code": code
        }

async def analyze_data_impl(data_code: str, analysis_code: str) -> Dict[str, Any]:
    """Data分析工具"""
    try:
        combined_code = f"{data_code}\n\n{analysis_code}"
        
        # 安全性检查
        is_safe, message = is_safe_code(combined_code)
        if not is_safe:
            return {
                "status": "error",
                "message": f"代码安全检查Failed: {message}",
                "code": combined_code
            }
        
        # 准备分析环境
        safe_globals = {
            '__builtins__': {name: getattr(__builtins__, name) for name in SAFE_BUILTINS if hasattr(__builtins__, name)}
        }
        
        # 导入Data分析库
        try:
            import numpy as np
            import pandas as pd
            import matplotlib.pyplot as plt
            
            safe_globals.update({
                'numpy': np,
                'np': np,
                'pandas': pd,
                'pd': pd,
                'plt': plt
            })
            
            # 尝试导入其他分析库
            try:
                import seaborn as sns
                safe_globals['seaborn'] = sns
                safe_globals['sns'] = sns
            except ImportError:
                pass
            
            try:
                import scipy.stats as stats
                safe_globals['stats'] = stats
            except ImportError:
                pass
            
        except ImportError:
            return {
                "status": "error",
                "message": "Data分析库未安装"
            }
        
        # 执行分析代码
        output_buffer = io.StringIO()
        
        with contextlib.redirect_stdout(output_buffer):
            with contextlib.redirect_stderr(output_buffer):
                exec(combined_code, safe_globals)
        
        output = output_buffer.getvalue()
        
        return {
            "status": "success",
            "output": output,
            "data_code": data_code,
            "analysis_code": analysis_code,
            "message": "Data分析完成"
        }
    
    except Exception as e:
        error_traceback = traceback.format_exc()
        return {
            "status": "error",
            "message": str(e),
            "traceback": error_traceback,
            "code": combined_code
        }

async def install_package_impl(package_name: str) -> Dict[str, Any]:
    """安装Python包 (仅限白名单)"""
    # 允许安装的包白名单
    allowed_packages = {
        'numpy', 'pandas', 'matplotlib', 'seaborn', 'scipy', 'sklearn',
        'requests', 'beautifulsoup4', 'lxml', 'pillow', 'opencv-python',
        'plotly', 'bokeh', 'altair', 'streamlit', 'dash'
    }
    
    if package_name not in allowed_packages:
        return {
            "status": "error",
            "message": f"不允许安装包: {package_name}",
            "allowed_packages": list(allowed_packages)
        }
    
    try:
        result = subprocess.run(
            [sys.executable, '-m', 'pip', 'install', package_name],
            capture_output=True,
            text=True,
            timeout=120
        )
        
        if result.returncode == 0:
            return {
                "status": "success",
                "message": f"Success安装包: {package_name}",
                "output": result.stdout
            }
        else:
            return {
                "status": "error",
                "message": f"安装Failed: {package_name}",
                "error": result.stderr
            }
    
    except subprocess.TimeoutExpired:
        return {
            "status": "error",
            "message": f"安装超时: {package_name}"
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

async def main():
    """Run MCP server"""
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())

if __name__ == "__main__":
    asyncio.run(main()) 