MCP 与零信任架构

# mcp_zero_trust_access_control.py
from typing import List, Dict, Any, Tuple
import logging
import json
from datetime import datetime

class ZeroTrustPolicy:
    def __init__(self, policy_id: str, name: str, description: str, 
                 condition: Dict, action: str):
        self.policy_id = policy_id
        self.name = name
        self.description = description
        self.condition = condition
        self.action = action  # allow, deny
        self.logger = logging.getLogger(f"mcp_zero_trust.policy.{policy_id}")
    
    def evaluate(self, subject: Dict, object_: Dict, action: str, 
                environment: Dict) -> bool:
        """评估策略"""
        self.logger.debug(f"Evaluating policy {self.policy_id} for {action} operation")
        
        # 检查策略是否适用于当前操作
        if "action" in self.condition and self.condition["action"] != action:
            return False
        
        # 评估主体条件
        if "subject" in self.condition:
            for key, value in self.condition["subject"].items():
                if key not in subject or subject[key] != value:
                    return False
        
        # 评估客体条件
        if "object" in self.condition:
            for key, value in self.condition["object"].items():
                if key not in object_ or object_[key] != value:
                    return False
        
        # 评估环境条件
        if "environment" in self.condition:
            for key, env_condition in self.condition["environment"].items():
                if key not in environment:
                    return False
                
                actual_value = environment[key]
                
                if isinstance(env_condition, dict):
                    # 处理比较操作
                    if "operator" in env_condition and "value" in env_condition:
                        operator = env_condition["operator"]
                        compare_value = env_condition["value"]
                        
                        if operator == "gt":
                            if not (actual_value > compare_value):
                                return False
                        elif operator == "lt":
                            if not (actual_value < compare_value):
                                return False
                        elif operator == "eq":
                            if actual_value != compare_value:
                                return False
                        elif operator == "ne":
                            if actual_value == compare_value:
                                return False
                        elif operator == "contains":
                            if compare_value not in actual_value:
                                return False
                        elif operator == "in":
                            if actual_value not in compare_value:
                                return False
                else:
                    # 直接比较
                    if actual_value != env_condition:
                        return False
        
        self.logger.info(f"Policy {self.policy_id} matched for {action} operation")
        return True
    
    def to_dict(self) -> Dict:
        """转换为字典"""
        return {
            "policy_id": self.policy_id,
            "name": self.name,
            "description": self.description,
            "condition": self.condition,
            "action": self.action
        }

class ZeroTrustAccessControl:
    def __init__(self):
        self.policies = {}
        self.logger = logging.getLogger("mcp_zero_trust.access_control")
    
    def add_policy(self, policy: ZeroTrustPolicy):
        """添加策略"""
        self.policies[policy.policy_id] = policy
        self.logger.info(f"Added policy: {policy.policy_id}")
    
    def remove_policy(self, policy_id: str):
        """移除策略"""
        if policy_id in self.policies:
            del self.policies[policy_id]
            self.logger.info(f"Removed policy: {policy_id}")
    
    def update_policy(self, policy: ZeroTrustPolicy):
        """更新策略"""
        self.policies[policy.policy_id] = policy
        self.logger.info(f"Updated policy: {policy.policy_id}")
    
    def evaluate_access(self, subject: Dict, object_: Dict, action: str, 
                       environment: Dict) -> Tuple[bool, List[str]]:
        """评估访问请求"""
        self.logger.info(f"Evaluating access: subject={subject.get('id')}, object={object_.get('id')}, action={action}")
        
        matched_policies = []
        allow = False
        
        # 评估所有策略
        for policy in self.policies.values():
            if policy.evaluate(subject, object_, action, environment):
                matched_policies.append(policy.policy_id)
                if policy.action == "allow":
                    allow = True
                elif policy.action == "deny":
                    allow = False
                    # 拒绝策略优先级高于允许策略
                    break
        
        self.logger.info(f"Access decision: {'allow' if allow else 'deny'}, matched_policies={matched_policies}")
        return allow, matched_policies
    
    def get_policies(self) -> List[Dict]:
        """获取所有策略"""
        return [policy.to_dict() for policy in self.policies.values()]

# mcp_microsegmentation.py
from typing import List, Dict, Tuple
import logging
from datetime import datetime

class SecurityDomain:
    def __init__(self, domain_id: str, name: str, description: str, 
                 components: List[str], parent_domain: str = None):
        self.domain_id = domain_id
        self.name = name
        self.description = description
        self.components = components
        self.parent_domain = parent_domain
        self.logger = logging.getLogger(f"mcp_microsegmentation.domain.{domain_id}")
    
    def add_component(self, component_id: str):
        """添加组件到安全域"""
        if component_id not in self.components:
            self.components.append(component_id)
            self.logger.info(f"Added component {component_id} to domain {self.domain_id}")
    
    def remove_component(self, component_id: str):
        """从安全域中移除组件"""
        if component_id in self.components:
            self.components.remove(component_id)
            self.logger.info(f"Removed component {component_id} from domain {self.domain_id}")
    
    def to_dict(self) -> Dict:
        """转换为字典"""
        return {
            "domain_id": self.domain_id,
            "name": self.name,
            "description": self.description,
            "components": self.components,
            "parent_domain": self.parent_domain
        }

class MicrosegmentationPolicy:
    def __init__(self, policy_id: str, name: str, description: str, 
                 source_domain: str, destination_domain: str, 
                 protocol: str, port: int, action: str):
        self.policy_id = policy_id
        self.name = name
        self.description = description
        self.source_domain = source_domain
        self.destination_domain = destination_domain
        self.protocol = protocol
        self.port = port
        self.action = action  # allow, deny
        self.logger = logging.getLogger(f"mcp_microsegmentation.policy.{policy_id}")
    
    def evaluate_traffic(self, source_domain: str, destination_domain: str, 
                        protocol: str, port: int) -> bool:
        """评估流量是否符合策略"""
        if self.source_domain != source_domain:
            return False
        
        if self.destination_domain != destination_domain:
            return False
        
        if self.protocol != protocol:
            return False
        
        if self.port != port:
            return False
        
        self.logger.info(f"Microsegmentation policy {self.policy_id} matched: {self.action}")
        return self.action == "allow"
    
    def to_dict(self) -> Dict:
        """转换为字典"""
        return {
            "policy_id": self.policy_id,
            "name": self.name,
            "description": self.description,
            "source_domain": self.source_domain,
            "destination_domain": self.destination_domain,
            "protocol": self.protocol,
            "port": self.port,
            "action": self.action
        }

class MCPMicrosegmentation:
    def __init__(self):
        self.domains = {}
        self.policies = {}
        self.component_domain_map = {}
        self.logger = logging.getLogger("mcp_microsegmentation")
    
    def add_domain(self, domain: SecurityDomain):
        """添加安全域"""
        self.domains[domain.domain_id] = domain
        # 更新组件域映射
        for component in domain.components:
            self.component_domain_map[component] = domain.domain_id
        self.logger.info(f"Added domain: {domain.domain_id}")
    
    def remove_domain(self, domain_id: str):
        """移除安全域"""
        if domain_id in self.domains:
            # 清除组件域映射
            for component in self.domains[domain_id].components:
                if component in self.component_domain_map:
                    del self.component_domain_map[component]
            del self.domains[domain_id]
            self.logger.info(f"Removed domain: {domain_id}")
    
    def add_policy(self, policy: MicrosegmentationPolicy):
        """添加微分段策略"""
        self.policies[policy.policy_id] = policy
        self.logger.info(f"Added microsegmentation policy: {policy.policy_id}")
    
    def remove_policy(self, policy_id: str):
        """移除微分段策略"""
        if policy_id in self.policies:
            del self.policies[policy_id]
            self.logger.info(f"Removed microsegmentation policy: {policy_id}")
    
    def get_component_domain(self, component_id: str) -> str:
        """获取组件所属的安全域"""
        return self.component_domain_map.get(component_id, "unknown")
    
    def evaluate_traffic(self, source_component: str, destination_component: str, 
                        protocol: str, port: int) -> bool:
        """评估组件间流量是否允许"""
        source_domain = self.get_component_domain(source_component)
        destination_domain = self.get_component_domain(destination_component)
        
        self.logger.info(f"Evaluating traffic: source={source_component}({source_domain}), destination={destination_component}({destination_domain}), protocol={protocol}, port={port}")
        
        # 默认拒绝所有流量
        allowed = False
        
        # 评估所有策略
        for policy in self.policies.values():
            if policy.evaluate_traffic(source_domain, destination_domain, protocol, port):
                allowed = True
                break
        
        self.logger.info(f"Traffic decision: {'allowed' if allowed else 'denied'}")
        return allowed
    
    def get_domains(self) -> List[Dict]:
        """获取所有安全域"""
        return [domain.to_dict() for domain in self.domains.values()]
    
    def get_policies(self) -> List[Dict]:
        """获取所有微分段策略"""
        return [policy.to_dict() for policy in self.policies.values()]

# mcp_continuous_verification.py
from typing import Dict, List, Tuple
import logging
from datetime import datetime, timedelta
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
import joblib

class TrustScoreCalculator:
    def __init__(self, model_path: str = None):
        # 创建或加载信任评估模型
        if model_path:
            self.model = joblib.load(model_path)
        else:
            self.model = Pipeline([
                ("scaler", StandardScaler()),
                ("clf", RandomForestClassifier(n_estimators=100, random_state=42))
            ])
        
        self.logger = logging.getLogger("mcp_trust_score")
    
    def calculate_trust_score(self, subject: Dict, context: Dict) -> float:
        """计算信任评分"""
        self.logger.info(f"Calculating trust score for subject: {subject.get('id')}")
        
        # 提取特征
        features = self._extract_features(subject, context)
        
        # 计算信任评分（0-100）
        trust_score = self._calculate_score(features)
        
        self.logger.info(f"Trust score for subject {subject.get('id')}: {trust_score}")
        return trust_score
    
    def _extract_features(self, subject: Dict, context: Dict) -> List[float]:
        """提取信任评估特征"""
        features = []
        
        # 身份可信度（0-1）
        identity_trust = subject.get("identity_trust", 0.5)
        features.append(identity_trust)
        
        # 设备安全分数（0-100）
        device_score = context.get("device_score", 50)
        features.append(device_score / 100)
        
        # 网络安全分数（0-100）
        network_score = context.get("network_score", 50)
        features.append(network_score / 100)
        
        # 最近登录失败次数（0-10）
        failed_logins = context.get("failed_logins", 0)
        features.append(min(failed_logins, 10) / 10)
        
        # 异常行为分数（0-1）
        anomaly_score = context.get("anomaly_score", 0.5)
        features.append(anomaly_score)
        
        # 上次密码更改天数（0-365）
        password_age_days = context.get("password_age_days", 0)
        features.append(min(password_age_days, 365) / 365)
        
        return features
    
    def _calculate_score(self, features: List[float]) -> float:
        """计算信任评分"""
        # 简化处理，实际应使用训练好的模型
        # 这里基于特征加权计算信任评分
        weights = [0.3, 0.2, 0.2, 0.15, 0.1, 0.05]
        
        trust_score = 0.0
        for feature, weight in zip(features, weights):
            trust_score += feature * weight
        
        # 转换为 0-100 分
        trust_score = trust_score * 100
        
        return max(0.0, min(100.0, trust_score))

class ContinuousVerificationSystem:
    def __init__(self, trust_score_threshold: float = 70.0):
        self.trust_calculator = TrustScoreCalculator()
        self.trust_score_threshold = trust_score_threshold
        self.session_trust_scores = {}
        self.logger = logging.getLogger("mcp_continuous_verification")
    
    def verify_session(self, session_id: str, subject: Dict, context: Dict) -> Tuple[bool, float]:
        """验证会话"""
        self.logger.info(f"Verifying session: {session_id}")
        
        # 计算当前信任评分
        trust_score = self.trust_calculator.calculate_trust_score(subject, context)
        
        # 更新会话信任评分
        self.session_trust_scores[session_id] = {
            "trust_score": trust_score,
            "timestamp": datetime.now()
        }
        
        # 判断是否允许继续访问
        allowed = trust_score >= self.trust_score_threshold
        
        self.logger.info(f"Session verification result: allowed={allowed}, trust_score={trust_score}")
        return allowed, trust_score
    
    def update_authorization(self, session_id: str, subject: Dict, context: Dict, 
                           current_authorization: Dict) -> Dict:
        """动态更新授权"""
        self.logger.info(f"Updating authorization for session: {session_id}")
        
        # 验证会话
        allowed, trust_score = self.verify_session(session_id, subject, context)
        
        # 如果不允许继续访问，撤销所有授权
        if not allowed:
            self.logger.warning(f"Revoking authorization for session {session_id} due to low trust score: {trust_score}")
            return {"allowed": False, "reason": "low_trust_score", "trust_score": trust_score}
        
        # 动态调整授权
        updated_authorization = current_authorization.copy()
        
        # 根据信任评分调整权限范围
        if trust_score >= 90.0:
            # 高信任评分，允许所有操作
            updated_authorization["permissions"] = "full"
        elif trust_score >= 70.0:
            # 中高信任评分，允许大部分操作
            updated_authorization["permissions"] = "partial"
        elif trust_score >= 50.0:
            # 中信任评分，允许基本操作
            updated_authorization["permissions"] = "basic"
        else:
            # 低信任评分，只允许最小操作
            updated_authorization["permissions"] = "minimal"
        
        # 添加信任评分信息
        updated_authorization["trust_score"] = trust_score
        updated_authorization["last_verified"] = datetime.now().isoformat()
        
        self.logger.info(f"Updated authorization for session {session_id}: permissions={updated_authorization['permissions']}")
        return updated_authorization
    
    def get_session_trust_history(self, session_id: str) -> List[Dict]:
        """获取会话信任评分历史"""
        # 简化处理，返回最近的信任评分
        if session_id in self.session_trust_scores:
            return [self.session_trust_scores[session_id]]
        return []
    
    def cleanup_old_sessions(self, max_age_hours: int = 24):
        """清理旧会话"""
        self.logger.info(f"Cleaning up sessions older than {max_age_hours} hours")
        
        cutoff_time = datetime.now() - timedelta(hours=max_age_hours)
        old_sessions = []
        
        for session_id, data in self.session_trust_scores.items():
            if data["timestamp"] < cutoff_time:
                old_sessions.append(session_id)
        
        for session_id in old_sessions:
            del self.session_trust_scores[session_id]
        
        self.logger.info(f"Cleaned up {len(old_sessions)} old sessions")

# mcp_zero_trust_integration.py
from typing import Dict, List, Tuple
import logging
from datetime import datetime

# 导入前面实现的零信任组件
from mcp_zero_trust_access_control import ZeroTrustAccessControl, ZeroTrustPolicy
from mcp_microsegmentation import MCPMicrosegmentation, SecurityDomain, MicrosegmentationPolicy
from mcp_continuous_verification import ContinuousVerificationSystem

class MCPZeroTrustArchitecture:
    def __init__(self):
        # 初始化零信任组件
        self.access_control = ZeroTrustAccessControl()
        self.microsegmentation = MCPMicrosegmentation()
        self.continuous_verification = ContinuousVerificationSystem()
        
        self.logger = logging.getLogger("mcp_zero_trust")
        self.logger.info("MCP Zero Trust Architecture initialized")
    
    def initialize_default_config(self):
        """初始化默认配置"""
        self.logger.info("Initializing default zero trust configuration")
        
        # 1. 初始化安全域
        self._init_default_domains()
        
        # 2. 初始化微分段策略
        self._init_default_microsegmentation_policies()
        
        # 3. 初始化访问控制策略
        self._init_default_access_policies()
        
        self.logger.info("Default zero trust configuration initialized successfully")
    
    def _init_default_domains(self):
        """初始化默认安全域"""
        # 客户端域
        client_domain = SecurityDomain(
            domain_id="client_domain",
            name="MCP 客户端域",
            description="MCP Client 所在的安全域",
            components=["client_001", "client_002", "client_003"]
        )
        self.microsegmentation.add_domain(client_domain)
        
        # 服务器域
        server_domain = SecurityDomain(
            domain_id="server_domain",
            name="MCP 服务器域",
            description="MCP Server 所在的安全域",
            components=["server_001", "server_002"]
        )
        self.microsegmentation.add_domain(server_domain)
        
        # 主机域
        host_domain = SecurityDomain(
            domain_id="host_domain",
            name="MCP 主机域",
            description="MCP Host 所在的安全域",
            components=["host_001", "host_002", "host_003"]
        )
        self.microsegmentation.add_domain(host_domain)
        
        # 工具域
        tool_domain = SecurityDomain(
            domain_id="tool_domain",
            name="MCP 工具域",
            description="MCP 工具所在的安全域",
            components=["tool_001", "tool_002", "tool_003"]
        )
        self.microsegmentation.add_domain(tool_domain)
    
    def _init_default_microsegmentation_policies(self):
        """初始化默认微分段策略"""
        # 允许客户端到服务器的 HTTPS 通信
        policy1 = MicrosegmentationPolicy(
            policy_id="policy_001",
            name="客户端到服务器 HTTPS 通信",
            description="允许 MCP 客户端域到服务器域的 HTTPS 通信",
            source_domain="client_domain",
            destination_domain="server_domain",
            protocol="TCP",
            port=443,
            action="allow"
        )
        self.microsegmentation.add_policy(policy1)
        
        # 允许服务器到主机的 gRPC 通信
        policy2 = MicrosegmentationPolicy(
            policy_id="policy_002",
            name="服务器到主机 gRPC 通信",
            description="允许 MCP 服务器域到主机域的 gRPC 通信",
            source_domain="server_domain",
            destination_domain="host_domain",
            protocol="TCP",
            port=50051,
            action="allow"
        )
        self.microsegmentation.add_policy(policy2)
        
        # 允许主机到工具的通信
        policy3 = MicrosegmentationPolicy(
            policy_id="policy_003",
            name="主机到工具通信",
            description="允许 MCP 主机域到工具域的通信",
            source_domain="host_domain",
            destination_domain="tool_domain",
            protocol="TCP",
            port=0,  # 允许所有端口
            action="allow"
        )
        self.microsegmentation.add_policy(policy3)
    
    def _init_default_access_policies(self):
        """初始化默认访问控制策略"""
        # 允许管理员执行所有操作
        policy1 = ZeroTrustPolicy(
            policy_id="access_policy_001",
            name="管理员全权限",
            description="允许管理员角色执行所有 MCP 操作",
            condition={
                "subject": {
                    "role": "admin"
                }
            },
            action="allow"
        )
        self.access_control.add_policy(policy1)
        
        # 允许普通用户执行基本工具调用
        policy2 = ZeroTrustPolicy(
            policy_id="access_policy_002",
            name="普通用户基本操作",
            description="允许普通用户角色执行基本 MCP 工具调用",
            condition={
                "subject": {
                    "role": "user"
                },
                "action": {
                    "in": ["tool_call", "status_check"]
                },
                "environment": {
                    "device_score": {
                        "operator": "gt",
                        "value": 70
                    }
                }
            },
            action="allow"
        )
        self.access_control.add_policy(policy2)
        
        # 拒绝未认证用户访问
        policy3 = ZeroTrustPolicy(
            policy_id="access_policy_003",
            name="拒绝未认证用户",
            description="拒绝未认证用户访问 MCP 系统",
            condition={
                "subject": {
                    "authenticated": False
                }
            },
            action="deny"
        )
        self.access_control.add_policy(policy3)
    
    def process_tool_call_request(self, client_id: str, tool_id: str, 
                                 parameters: Dict, context: Dict) -> Dict:
        """处理 MCP 工具调用请求"""
        self.logger.info(f"Processing tool call request: client_id={client_id}, tool_id={tool_id}")
        
        # 1. 微分段检查：检查客户端到服务器的流量是否允许
        microsegmentation_allowed = self.microsegmentation.evaluate_traffic(
            source_component=client_id,
            destination_component="server_001",
            protocol="TCP",
            port=443
        )
        
        if not microsegmentation_allowed:
            return {
                "status": "error",
                "code": "MICROSEGMENTATION_DENIED",
                "message": "流量被微分段策略拒绝"
            }
        
        # 2. 访问控制检查：评估访问权限
        subject = {
            "id": client_id,
            "role": context.get("user_role", "user"),
            "authenticated": context.get("authenticated", True)
        }
        
        object_ = {
            "id": tool_id,
            "type": "tool"
        }
        
        action = "tool_call"
        
        access_allowed, matched_policies = self.access_control.evaluate_access(
            subject=subject,
            object_=object_,
            action=action,
            environment=context
        )
        
        if not access_allowed:
            return {
                "status": "error",
                "code": "ACCESS_DENIED",
                "message": "访问被拒绝"
            }
        
        # 3. 信任评估与持续验证
        trust_score = self.continuous_verification.trust_calculator.calculate_trust_score(
            subject=subject,
            context=context
        )
        
        # 4. 动态授权
        initial_authorization = {
            "allowed": True,
            "permissions": "partial",
            "tool_id": tool_id,
            "client_id": client_id
        }
        
        session_id = context.get("session_id", f"session_{client_id}_{datetime.now().timestamp()}")
        updated_authorization = self.continuous_verification.update_authorization(
            session_id=session_id,
            subject=subject,
            context=context,
            current_authorization=initial_authorization
        )
        
        # 5. 生成响应
        response = {
            "status": "success" if updated_authorization["allowed"] else "error",
            "code": "OK" if updated_authorization["allowed"] else "AUTHORIZATION_REVOKED",
            "message": "工具调用授权成功" if updated_authorization["allowed"] else "授权被撤销",
            "authorization": updated_authorization,
            "trust_score": trust_score,
            "matched_policies": matched_policies
        }
        
        self.logger.info(f"Tool call request processed: status={response['status']}, trust_score={trust_score}")
        return response
    
    def get_zero_trust_status(self) -> Dict:
        """获取零信任架构状态"""
        status = {
            "timestamp": datetime.now().isoformat(),
            "domains": self.microsegmentation.get_domains(),
            "microsegmentation_policies": self.microsegmentation.get_policies(),
            "access_policies": self.access_control.get_policies()
        }
        
        return status

# 示例使用
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建 MCP 零信任架构实例
    zta = MCPZeroTrustArchitecture()
    
    # 初始化默认配置
    zta.initialize_default_config()
    
    # 模拟工具调用请求
    client_id = "client_001"
    tool_id = "tool_001"
    parameters = {
        "file_path": "/data/test.txt",
        "mode": "read"
    }
    
    # 上下文信息
    context = {
        "user_role": "user",
        "authenticated": True,
        "device_score": 85,
        "network_score": 90,
        "failed_logins": 0,
        "anomaly_score": 0.1,
        "password_age_days": 30,
        "session_id": "session_001"
    }
    
    # 处理工具调用请求
    response = zta.process_tool_call_request(
        client_id=client_id,
        tool_id=tool_id,
        parameters=parameters,
        context=context
    )
    
    print(f"Tool call response: {response}")
    
    # 获取零信任架构状态
    zta_status = zta.get_zero_trust_status()
    print(f"Zero Trust Architecture Status: {zta_status}")

# 安装依赖
pip install fastapi uvicorn sqlalchemy pydantic redis kafka-python numpy scikit-learn

# 配置零信任架构
cp config.example.yaml config.yaml
# 编辑配置文件
vim config.yaml

# 启动基础设施服务
docker-compose up -d postgresql redis kafka zookeeper

# 初始化数据库
python init_db.py --config config.yaml

# 启动零信任服务
python mcp_zero_trust_service.py --config config.yaml

# 启动微分段服务
python mcp_microsegmentation_service.py --config config.yaml

# 启动持续验证服务
python mcp_continuous_verification_service.py --config config.yaml