MCP 的白名单与沙箱机制

# mcp_whitelist_manager.py
from typing import List, Dict, Optional
from datetime import datetime
import json
import hashlib
import logging

class WhitelistItem:
    def __init__(self, item_id: str, item_type: str, name: str, version: str, 
                 parameters: Dict, enabled: bool, create_time: datetime, 
                 update_time: datetime, description: str = ""):
        self.item_id = item_id
        self.item_type = item_type  # global, system, tool, version, parameter
        self.name = name
        self.version = version
        self.parameters = parameters
        self.enabled = enabled
        self.create_time = create_time
        self.update_time = update_time
        self.description = description
    
    def to_dict(self) -> Dict:
        return {
            "item_id": self.item_id,
            "item_type": self.item_type,
            "name": self.name,
            "version": self.version,
            "parameters": self.parameters,
            "enabled": self.enabled,
            "create_time": self.create_time.isoformat(),
            "update_time": self.update_time.isoformat(),
            "description": self.description
        }
    
    def generate_hash(self) -> str:
        """生成白名单项的哈希值，用于检测变更"""
        data = json.dumps(self.to_dict(), sort_keys=True)
        return hashlib.sha256(data.encode()).hexdigest()

class WhitelistManager:
    def __init__(self, db_uri: str):
        self.db_uri = db_uri
        self.whitelist_items = {}
        self.logger = logging.getLogger("mcp_whitelist")
        self.load_whitelist()
    
    def load_whitelist(self):
        """从数据库加载白名单"""
        # 实现从数据库加载白名单的逻辑
        pass
    
    def save_whitelist(self):
        """保存白名单到数据库"""
        # 实现保存白名单到数据库的逻辑
        pass
    
    def create_whitelist_item(self, item_type: str, name: str, version: str, 
                             parameters: Dict, enabled: bool = True, 
                             description: str = "") -> WhitelistItem:
        """创建白名单项"""
        item_id = f"{item_type}_{name}_{version}_{datetime.now().timestamp()}"
        now = datetime.now()
        item = WhitelistItem(
            item_id=item_id,
            item_type=item_type,
            name=name,
            version=version,
            parameters=parameters,
            enabled=enabled,
            create_time=now,
            update_time=now,
            description=description
        )
        
        self.whitelist_items[item_id] = item
        self.save_whitelist()
        self.logger.info(f"Created whitelist item: {item_id}")
        return item
    
    def update_whitelist_item(self, item_id: str, **kwargs) -> Optional[WhitelistItem]:
        """更新白名单项"""
        if item_id not in self.whitelist_items:
            return None
        
        item = self.whitelist_items[item_id]
        for key, value in kwargs.items():
            if hasattr(item, key):
                setattr(item, key, value)
        
        item.update_time = datetime.now()
        self.save_whitelist()
        self.logger.info(f"Updated whitelist item: {item_id}")
        return item
    
    def delete_whitelist_item(self, item_id: str) -> bool:
        """删除白名单项"""
        if item_id in self.whitelist_items:
            del self.whitelist_items[item_id]
            self.save_whitelist()
            self.logger.info(f"Deleted whitelist item: {item_id}")
            return True
        return False
    
    def get_whitelist_item(self, item_id: str) -> Optional[WhitelistItem]:
        """获取白名单项"""
        return self.whitelist_items.get(item_id)
    
    def list_whitelist_items(self, item_type: Optional[str] = None, 
                            enabled: Optional[bool] = None) -> List[WhitelistItem]:
        """列出白名单项"""
        items = []
        for item in self.whitelist_items.values():
            if item_type and item.item_type != item_type:
                continue
            if enabled is not None and item.enabled != enabled:
                continue
            items.append(item)
        return items
    
    def check_whitelist(self, item_type: str, name: str, version: str, 
                       parameters: Dict) -> bool:
        """检查是否在白名单中"""
        # 1. 检查全局白名单
        global_items = self.list_whitelist_items(item_type="global", enabled=True)
        if not global_items:
            return False
        
        # 2. 检查系统白名单
        system_items = self.list_whitelist_items(item_type="system", enabled=True)
        if not system_items:
            return False
        
        # 3. 检查工具白名单
        tool_items = self.list_whitelist_items(item_type="tool", enabled=True)
        tool_allowed = False
        for item in tool_items:
            if item.name == name:
                tool_allowed = True
                break
        if not tool_allowed:
            return False
        
        # 4. 检查版本白名单
        version_items = self.list_whitelist_items(item_type="version", enabled=True)
        version_allowed = False
        for item in version_items:
            if item.name == name and self._check_version(item.version, version):
                version_allowed = True
                break
        if not version_allowed:
            return False
        
        # 5. 检查参数白名单
        parameter_items = self.list_whitelist_items(item_type="parameter", enabled=True)
        for item in parameter_items:
            if item.name == name:
                for param_name, param_value in item.parameters.items():
                    if param_name in parameters:
                        if not self._check_parameter(param_value, parameters[param_name]):
                            return False
        
        return True
    
    def _check_version(self, allowed_version: str, actual_version: str) -> bool:
        """检查版本是否允许"""
        # 实现版本检查逻辑，支持精确匹配、范围匹配等
        if allowed_version == "*":
            return True
        return allowed_version == actual_version
    
    def _check_parameter(self, allowed_value: str, actual_value: str) -> bool:
        """检查参数是否允许"""
        # 实现参数检查逻辑，支持精确匹配、正则表达式匹配等
        if allowed_value == "*":
            return True
        return allowed_value == actual_value
    
    def push_update(self, mcp_server: str):
        """向 MCP Server 推送白名单更新"""
        # 实现白名单更新推送逻辑
        pass
    
    def record_audit_log(self, action: str, item_id: str, operator: str, 
                        details: Dict):
        """记录白名单变更审计日志"""
        # 实现审计日志记录逻辑
        pass

# mcp_dynamic_whitelist.py
from typing import List, Dict
from datetime import datetime
import machine_learning as ml
from mcp_whitelist_manager import WhitelistManager

class DynamicWhitelistSystem:
    def __init__(self, whitelist_manager: WhitelistManager):
        self.whitelist_manager = whitelist_manager
        self.ml_model = ml.load_model("whitelist_model")
        self.audit_logs = []
    
    def learn_from_audit_logs(self, logs: List[Dict]):
        """从审计日志中学习白名单模式"""
        # 使用机器学习模型分析审计日志，学习正常的工具调用模式
        features = self._extract_features(logs)
        self.ml_model.train(features)
        
        # 根据学习结果优化白名单配置
        optimized_whitelist = self.ml_model.generate_whitelist()
        for item in optimized_whitelist:
            self.whitelist_manager.create_whitelist_item(**item)
    
    def dynamic_adjust_whitelist(self, context: Dict):
        """根据上下文动态调整白名单"""
        # 分析当前上下文，预测需要的白名单配置
        prediction = self.ml_model.predict(context)
        
        # 根据预测结果调整白名单
        for item in prediction["add_items"]:
            self.whitelist_manager.create_whitelist_item(**item)
        
        for item_id in prediction["remove_items"]:
            self.whitelist_manager.delete_whitelist_item(item_id)
        
        for item_id, updates in prediction["update_items"].items():
            self.whitelist_manager.update_whitelist_item(item_id, **updates)
    
    def _extract_features(self, logs: List[Dict]) -> List[Dict]:
        """从审计日志中提取特征"""
        features = []
        for log in logs:
            feature = {
                "tool_name": log.get("tool_name"),
                "tool_version": log.get("tool_version"),
                "parameters": log.get("parameters"),
                "timestamp": log.get("timestamp"),
                "user_id": log.get("user_id"),
                "client_id": log.get("client_id"),
                "server_id": log.get("server_id"),
                "result": log.get("result"),
                "duration_ms": log.get("duration_ms")
            }
            features.append(feature)
        return features
    
    def evaluate_whitelist_effectiveness(self) -> Dict:
        """评估白名单有效性"""
        # 分析白名单的误判率、漏判率等指标
        metrics = {
            "true_positive": 0,
            "false_positive": 0,
            "true_negative": 0,
            "false_negative": 0
        }
        
        # 实现评估逻辑
        
        return metrics

# mcp_sandbox_manager.py
from typing import List, Dict, Optional
import docker
import logging
import uuid
import time

class SandboxConfig:
    def __init__(self, cpu_limit: float = 1.0, memory_limit: str = "128m", 
                 disk_limit: str = "1g", network_mode: str = "bridge", 
                 ports: Dict = None, volumes: Dict = None, 
                 environment: Dict = None):
        self.cpu_limit = cpu_limit
        self.memory_limit = memory_limit
        self.disk_limit = disk_limit
        self.network_mode = network_mode
        self.ports = ports or {}
        self.volumes = volumes or {}
        self.environment = environment or {}
    
    def to_dict(self) -> Dict:
        return {
            "cpu_limit": self.cpu_limit,
            "memory_limit": self.memory_limit,
            "disk_limit": self.disk_limit,
            "network_mode": self.network_mode,
            "ports": self.ports,
            "volumes": self.volumes,
            "environment": self.environment
        }

class Sandbox:
    def __init__(self, sandbox_id: str, container_id: str, image: str, 
                 config: SandboxConfig, status: str = "created", 
                 create_time: float = 0):
        self.sandbox_id = sandbox_id
        self.container_id = container_id
        self.image = image
        self.config = config
        self.status = status
        self.create_time = create_time
    
    def to_dict(self) -> Dict:
        return {
            "sandbox_id": self.sandbox_id,
            "container_id": self.container_id,
            "image": self.image,
            "config": self.config.to_dict(),
            "status": self.status,
            "create_time": self.create_time
        }

class SandboxManager:
    def __init__(self, docker_client: docker.DockerClient = None):
        self.docker_client = docker_client or docker.from_env()
        self.sandboxes = {}
        self.logger = logging.getLogger("mcp_sandbox")
    
    def create_sandbox(self, image: str, config: SandboxConfig, 
                      command: str = None, entrypoint: str = None) -> Sandbox:
        """创建沙箱"""
        sandbox_id = str(uuid.uuid4())
        create_time = time.time()
        
        try:
            # 创建容器
            container = self.docker_client.containers.create(
                image=image,
                command=command,
                entrypoint=entrypoint,
                cpu_quota=int(config.cpu_limit * 100000),
                mem_limit=config.memory_limit,
                network_mode=config.network_mode,
                ports=config.ports,
                volumes=config.volumes,
                environment=config.environment,
                detach=True
            )
            
            # 启动容器
            container.start()
            
            sandbox = Sandbox(
                sandbox_id=sandbox_id,
                container_id=container.id,
                image=image,
                config=config,
                status="running",
                create_time=create_time
            )
            
            self.sandboxes[sandbox_id] = sandbox
            self.logger.info(f"Created sandbox: {sandbox_id}, container: {container.id}")
            return sandbox
        except Exception as e:
            self.logger.error(f"Failed to create sandbox: {e}")
            raise
    
    def destroy_sandbox(self, sandbox_id: str) -> bool:
        """销毁沙箱"""
        if sandbox_id not in self.sandboxes:
            return False
        
        sandbox = self.sandboxes[sandbox_id]
        try:
            # 停止容器
            container = self.docker_client.containers.get(sandbox.container_id)
            container.stop(timeout=5)
            
            # 删除容器
            container.remove(force=True)
            
            del self.sandboxes[sandbox_id]
            self.logger.info(f"Destroyed sandbox: {sandbox_id}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to destroy sandbox: {e}")
            return False
    
    def get_sandbox(self, sandbox_id: str) -> Optional[Sandbox]:
        """获取沙箱信息"""
        return self.sandboxes.get(sandbox_id)
    
    def list_sandboxes(self, status: Optional[str] = None) -> List[Sandbox]:
        """列出沙箱"""
        sandboxes = []
        for sandbox in self.sandboxes.values():
            if status and sandbox.status != status:
                continue
            sandboxes.append(sandbox)
        return sandboxes
    
    def execute_in_sandbox(self, sandbox_id: str, command: List[str], 
                          timeout: int = 30) -> Dict:
        """在沙箱中执行命令"""
        if sandbox_id not in self.sandboxes:
            return {"success": False, "error": "Sandbox not found"}
        
        sandbox = self.sandboxes[sandbox_id]
        try:
            container = self.docker_client.containers.get(sandbox.container_id)
            result = container.exec_run(
                cmd=command,
                timeout=timeout,
                stdout=True,
                stderr=True
            )
            
            return {
                "success": True,
                "exit_code": result.exit_code,
                "stdout": result.output.decode("utf-8"),
                "stderr": result.output.decode("utf-8") if result.exit_code != 0 else ""
            }
        except Exception as e:
            self.logger.error(f"Failed to execute in sandbox: {e}")
            return {"success": False, "error": str(e)}
    
    def get_sandbox_stats(self, sandbox_id: str) -> Dict:
        """获取沙箱资源使用统计"""
        if sandbox_id not in self.sandboxes:
            return {"success": False, "error": "Sandbox not found"}
        
        sandbox = self.sandboxes[sandbox_id]
        try:
            container = self.docker_client.containers.get(sandbox.container_id)
            stats = container.stats(stream=False)
            
            return {
                "success": True,
                "stats": {
                    "cpu_usage": stats["cpu_stats"]["cpu_usage"]["total_usage"],
                    "memory_usage": stats["memory_stats"]["usage"],
                    "memory_limit": stats["memory_stats"]["limit"],
                    "network_io": {
                        "rx_bytes": stats["networks"]["eth0"]["rx_bytes"],
                        "tx_bytes": stats["networks"]["eth0"]["tx_bytes"]
                    },
                    "block_io": {
                        "read_bytes": stats["blkio_stats"]["io_service_bytes_recursive"][0]["value"] if stats["blkio_stats"]["io_service_bytes_recursive"] else 0,
                        "write_bytes": stats["blkio_stats"]["io_service_bytes_recursive"][1]["value"] if len(stats["blkio_stats"]["io_service_bytes_recursive"]) > 1 else 0
                    }
                }
            }
        except Exception as e:
            self.logger.error(f"Failed to get sandbox stats: {e}")
            return {"success": False, "error": str(e)}
    
    def update_sandbox_config(self, sandbox_id: str, config: SandboxConfig) -> bool:
        """更新沙箱配置"""
        if sandbox_id not in self.sandboxes:
            return False
        
        sandbox = self.sandboxes[sandbox_id]
        try:
            container = self.docker_client.containers.get(sandbox.container_id)
            
            # 更新容器资源限制
            container.update(
                cpu_quota=int(config.cpu_limit * 100000),
                mem_limit=config.memory_limit
            )
            
            sandbox.config = config
            self.logger.info(f"Updated sandbox config: {sandbox_id}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to update sandbox config: {e}")
            return False
    
    def monitor_sandbox(self, sandbox_id: str, metrics: List[str] = None) -> Dict:
        """监控沙箱"""
        # 实现沙箱监控逻辑
        pass
    
    def clean_stale_sandboxes(self, max_age_seconds: int = 3600) -> int:
        """清理过期沙箱"""
        cleaned_count = 0
        current_time = time.time()
        
        for sandbox_id, sandbox in list(self.sandboxes.items()):
            if current_time - sandbox.create_time > max_age_seconds:
                if self.destroy_sandbox(sandbox_id):
                    cleaned_count += 1
        
        self.logger.info(f"Cleaned {cleaned_count} stale sandboxes")
        return cleaned_count

# mcp_sandbox_escape_detection.py
from typing import List, Dict
import logging
import time
from mcp_sandbox_manager import SandboxManager

class EscapeDetectionRule:
    def __init__(self, rule_id: str, name: str, description: str, 
                 severity: str, condition: Dict, action: str):
        self.rule_id = rule_id
        self.name = name
        self.description = description
        self.severity = severity  # low, medium, high, critical
        self.condition = condition
        self.action = action  # alert, destroy, quarantine
    
    def to_dict(self) -> Dict:
        return {
            "rule_id": self.rule_id,
            "name": self.name,
            "description": self.description,
            "severity": self.severity,
            "condition": self.condition,
            "action": self.action
        }

class EscapeDetectionEvent:
    def __init__(self, event_id: str, sandbox_id: str, rule_id: str, 
                 timestamp: float, details: Dict, severity: str):
        self.event_id = event_id
        self.sandbox_id = sandbox_id
        self.rule_id = rule_id
        self.timestamp = timestamp
        self.details = details
        self.severity = severity
    
    def to_dict(self) -> Dict:
        return {
            "event_id": self.event_id,
            "sandbox_id": self.sandbox_id,
            "rule_id": self.rule_id,
            "timestamp": self.timestamp,
            "details": self.details,
            "severity": self.severity
        }

class SandboxEscapeDetectionSystem:
    def __init__(self, sandbox_manager: SandboxManager):
        self.sandbox_manager = sandbox_manager
        self.rules = {}
        self.events = []
        self.logger = logging.getLogger("mcp_escape_detection")
    
    def add_detection_rule(self, rule: EscapeDetectionRule):
        """添加检测规则"""
        self.rules[rule.rule_id] = rule
        self.logger.info(f"Added detection rule: {rule.rule_id}")
    
    def remove_detection_rule(self, rule_id: str):
        """移除检测规则"""
        if rule_id in self.rules:
            del self.rules[rule_id]
            self.logger.info(f"Removed detection rule: {rule_id}")
    
    def detect_escape(self, sandbox_id: str) -> List[EscapeDetectionEvent]:
        """检测沙箱逃逸"""
        events = []
        sandbox = self.sandbox_manager.get_sandbox(sandbox_id)
        
        if not sandbox:
            return events
        
        # 获取沙箱统计信息
        stats = self.sandbox_manager.get_sandbox_stats(sandbox_id)
        if not stats["success"]:
            return events
        
        # 检查每个规则
        for rule in self.rules.values():
            if self._check_rule(rule, sandbox, stats["stats"]):
                # 检测到逃逸
                event = EscapeDetectionEvent(
                    event_id=f"escape_{sandbox_id}_{rule.rule_id}_{time.time()}",
                    sandbox_id=sandbox_id,
                    rule_id=rule.rule_id,
                    timestamp=time.time(),
                    details={"stats": stats["stats"]}, 
                    severity=rule.severity
                )
                events.append(event)
                self.events.append(event)
                
                # 执行响应动作
                self._execute_action(rule.action, sandbox_id)
        
        return events
    
    def _check_rule(self, rule: EscapeDetectionRule, sandbox: Sandbox, 
                   stats: Dict) -> bool:
        """检查规则是否触发"""
        # 实现规则检查逻辑
        condition = rule.condition
        
        if "cpu_usage" in condition:
            cpu_threshold = condition["cpu_usage"]
            if stats["cpu_usage"] > cpu_threshold:
                return True
        
        if "memory_usage" in condition:
            memory_threshold = condition["memory_usage"]
            if stats["memory_usage"] > memory_threshold:
                return True
        
        # 检查网络异常
        if "network_anomaly" in condition and condition["network_anomaly"]:
            # 实现网络异常检测逻辑
            pass
        
        # 检查文件系统异常
        if "filesystem_anomaly" in condition and condition["filesystem_anomaly"]:
            # 实现文件系统异常检测逻辑
            pass
        
        return False
    
    def _execute_action(self, action: str, sandbox_id: str):
        """执行响应动作"""
        self.logger.info(f"Executing action {action} for sandbox {sandbox_id}")
        
        if action == "alert":
            # 发送告警
            self._send_alert(sandbox_id)
        elif action == "destroy":
            # 销毁沙箱
            self.sandbox_manager.destroy_sandbox(sandbox_id)
        elif action == "quarantine":
            # 隔离沙箱
            self._quarantine_sandbox(sandbox_id)
    
    def _send_alert(self, sandbox_id: str):
        """发送告警"""
        # 实现告警发送逻辑
        self.logger.warning(f"Alert: Possible sandbox escape detected for {sandbox_id}")
    
    def _quarantine_sandbox(self, sandbox_id: str):
        """隔离沙箱"""
        # 实现沙箱隔离逻辑
        self.logger.info(f"Quarantined sandbox: {sandbox_id}")
    
    def start_monitoring(self, interval_seconds: int = 10):
        """启动监控"""
        self.logger.info(f"Starting sandbox escape monitoring with interval {interval_seconds}s")
        
        while True:
            for sandbox_id in list(self.sandbox_manager.sandboxes.keys()):
                self.detect_escape(sandbox_id)
            time.sleep(interval_seconds)
    
    def get_detection_events(self, sandbox_id: str = None, 
                           time_range: Dict = None) -> List[EscapeDetectionEvent]:
        """获取检测事件"""
        events = []
        for event in self.events:
            if sandbox_id and event.sandbox_id != sandbox_id:
                continue
            if time_range:
                start_time = time_range.get("start_time")
                end_time = time_range.get("end_time")
                if start_time and event.timestamp < start_time:
                    continue
                if end_time and event.timestamp > end_time:
                    continue
            events.append(event)
        return events

# 安装依赖
pip install mcp-whitelist-sandbox

# 配置白名单与沙箱系统
cp config.example.yaml config.yaml
# 编辑配置文件
vim config.yaml

# 启动白名单管理服务
mcp-whitelist-service --config config.yaml

# 启动沙箱管理服务
mcp-sandbox-service --config config.yaml

# 启动逃逸检测服务
mcp-escape-detection-service --config config.yaml