learn-spider/app.py

import os
from fastmcp import FastMCP
from spider.mail_qq import start as start_mail_qq_spider
import asyncio
import json
from typing import Optional
import logging
import db.sqlite as db


logging.basicConfig(level=logging.INFO)
# 日志会输出到 stderr，不会污染 stdout
logging.info("服务器启动")
mcp = FastMCP("spider-server")

SERVER_HOST = "https://10.10.40.19"

SHELL_DIR = "./bin"
@mcp.tool()
async def mail_qq_spider(account: str, instance_id: int) -> str:
    """qq邮箱爬虫,第一个参数是用户名
    Args:
        account: qq邮箱用户名
        instance_id: 实例编号
    Returns:
        str: 爬虫结果
    """
    data={}
    try:
        if not account:
            raise ValueError("账号不能为空")
        if not instance_id:
            raise ValueError("实例编号不能为空")
        cdp_port = 9223 + instance_id
        logging.info(f"启动QQ邮箱爬虫，账号: {account}, 实例编号: {instance_id}, CDP端口: {cdp_port}")
        result = await asyncio.to_thread(start_mail_qq_spider, account,cdp_port)
        data["status"] = 200
        data["result"] = result
        return json.dumps(data, ensure_ascii=False)
    except Exception as e:
        data["status"] = 500
        data["message"] = str(e)
        return json.dumps(data, ensure_ascii=False)
@mcp.tool()
async def start_vnc_server() -> str:
    """启动 VNC + Chrome 实例"""
    data = {
        "status": 200,
        "action": "start"
    }
    instance_id = db.generate_vnc_instance_id(max=50)
    try:
        vnc_port = 5900 + instance_id
        websockify_port = 6080 + instance_id
        cdp_port = 9223 + instance_id
        
        process = await asyncio.create_subprocess_exec(
            "docker", "run", "-d",
            "--rm",
            "--name", f"vnc-server-{instance_id}",
            "-p", f"{cdp_port}:9223",
            "-p", f"{vnc_port}:5900",
            "-p", f"{websockify_port}:6080",
            "-e", f"CDP_PORT={cdp_port}",
            "--shm-size=2gb",
            "vnc-server:latest",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await process.communicate()
        
        if process.returncode == 0:
            data["instance_id"] = instance_id
            data["message"] = f"Container vnc-browser-{instance_id} started"
            data["vnc_url"] = f"{SERVER_HOST}:{vnc_port}"
            data["websockify_url"] = f"{SERVER_HOST}:{websockify_port}/vnc_lite.html"
            data["cdp_url"] = f"{SERVER_HOST}:{cdp_port}"
            db.insert_instance_id(instance_id,vnc_port,cdp_port,websockify_port)
            logging.info(f"VNC实例启动成功: {data}") 
        else:
            data["status"] = 500
            data["message"] = stderr.decode() or stdout.decode() or "Unknown error"
            logging.error(f"VNC实例启动失败: {data}")
        
        return json.dumps(data, ensure_ascii=False)
        
    except Exception as e:
        data["status"] = 500
        data["message"] = str(e)
        return json.dumps(data, ensure_ascii=False)
@mcp.tool()
async def stop_vnc_server(instance_id) -> str:
    """停止指定的 VNC 实例
    
    Args:
        instance_id: 实例编号
    
    Returns:
        操作结果
    """
    data = {
        "instance_id": instance_id,
        "status": 200,
        "action": "stop",
    }
    
    try:
        container_name = f"vnc-browser-{instance_id}"
        
        # 停止容器
        process = await asyncio.create_subprocess_exec(
            "docker", "stop", container_name,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await process.communicate()
        
        if process.returncode == 0:
                # 删除容器
            await asyncio.create_subprocess_exec(
                "docker", "rm", container_name,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            data["message"] = f"Container {container_name} stopped"
            db.delete_instance_id(instance_id)
            logging.info(f"VNC实例停止成功: {data}")
        else:
            data["status"] = 500
            data["message"] = stderr.decode() or stdout.decode() or "Unknown error"
            logging.error(f"VNC实例停止失败: {data}")
        return json.dumps(data, ensure_ascii=False)
        
    except Exception as e:
        data["status"] = 500
        data["message"] = str(e)
        return json.dumps(data, ensure_ascii=False)

# @mcp.tool()
async def start_spider(spider_code: str) -> str:
    """启动爬虫
    
    Args:
        spider_code: 爬虫代码
        display: 显示桌面编号
    """
    data={"display": display}
    try:
        # 调用外部 start_spider.sh 脚本，传入 spider_code 作为参数
        process = await asyncio.create_subprocess_exec(
            f"{SHELL_DIR}/start_spider.sh",
            spider_code,
            display,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()
        data["message"] = stdout.decode()
        data["status"] = 200 if process.returncode == 0 else 500
        if process.returncode == 0:
            return json.dumps(data, ensure_ascii=False)
        else:
            return json.dumps(data, ensure_ascii=False)
    except Exception as e:
        data["status"] = 500
        return json.dumps(data, ensure_ascii=False)


@mcp.tool()
def get_vnc_instances() -> str:
    """ 获取所有vnc服务器"""
    servers_db =  db.select_list()
    data = {
        "status": 200,
        "data": json.dumps(servers_db, ensure_ascii=False)
    }
    return json.dumps(data, ensure_ascii=False)

@mcp.resource("data://spider_code")
def get_spider_code_list() -> str:
    spider_list = [
        {"code": "01","name": "阳光采购爬虫","auth_name":"mail"},
        {"code": "02","name": "QQ邮箱爬虫","auth_name":"mail"},
    ]
    return json.dumps(spider_list, ensure_ascii=False)


# 运行服务器
if __name__ == "__main__":
    mcp.run(transport="sse")
    # mcp.run(transport="stdio")