MCP 的安全威胁模型

import json
import hashlib
from typing import Dict, List, Any

class MCPThreatDetector:
    def __init__(self):
        # 初始化威胁检测规则
        self.rules = {
            "spoofing": [
                {"name": "invalid_signature", "description": "无效的数字签名", "severity": "high"},
                {"name": "unknown_client", "description": "未知的客户端身份", "severity": "medium"}
            ],
            "tampering": [
                {"name": "modified_parameters", "description": "修改的调用参数", "severity": "medium"},
                {"name": "modified_tool_definition", "description": "修改的工具定义", "severity": "high"}
            ],
            "information_disclosure": [
                {"name": "sensitive_data_leak", "description": "敏感数据泄露", "severity": "high"},
                {"name": "unencrypted_transmission", "description": "未加密传输", "severity": "medium"}
            ],
            "denial_of_service": [
                {"name": "high_frequency_calls", "description": "高频调用", "severity": "medium"},
                {"name": "resource_exhaustion", "description": "资源耗尽", "severity": "high"}
            ],
            "elevation_of_privilege": [
                {"name": "privilege_escalation", "description": "权限提升", "severity": "high"},
                {"name": "unauthorized_access", "description": "未授权访问", "severity": "high"}
            ]
        }
        
        # 敏感数据模式
        self.sensitive_patterns = [
            r"password", r"secret", r"token", r"api_key",
            r"credit_card", r"social_security", r"personal_id"
        ]
    
    def detect_spoofing(self, request: Dict[str, Any]) -> List[Dict[str, Any]]:
        """检测身份伪造威胁"""
        threats = []
        
        # 检查签名
        if "signature" not in request or not self.verify_signature(request):
            threats.append({
                "type": "spoofing",
                "name": "invalid_signature",
                "description": "无效的数字签名",
                "severity": "high",
                "evidence": f"Missing or invalid signature for request {request.get('request_id')}"
            })
        
        # 检查客户端身份
        if "client_id" not in request or request["client_id"] not in self.get_valid_clients():
            threats.append({
                "type": "spoofing",
                "name": "unknown_client",
                "description": "未知的客户端身份",
                "severity": "medium",
                "evidence": f"Unknown client {request.get('client_id')}"
            })
        
        return threats
    
    def detect_tampering(self, request: Dict[str, Any]) -> List[Dict[str, Any]]:
        """检测篡改威胁"""
        threats = []
        
        # 检查参数完整性
        if "parameters_hash" in request:
            calculated_hash = self.calculate_parameters_hash(request.get("parameters", {}))
            if calculated_hash != request["parameters_hash"]:
                threats.append({
                    "type": "tampering",
                    "name": "modified_parameters",
                    "description": "修改的调用参数",
                    "severity": "medium",
                    "evidence": f"Parameters hash mismatch: expected {request['parameters_hash']}, got {calculated_hash}"
                })
        
        return threats
    
    def detect_information_disclosure(self, response: Dict[str, Any]) -> List[Dict[str, Any]]:
        """检测信息泄露威胁"""
        threats = []
        
        # 检查敏感数据泄露
        response_str = json.dumps(response)
        for pattern in self.sensitive_patterns:
            if pattern in response_str.lower():
                threats.append({
                    "type": "information_disclosure",
                    "name": "sensitive_data_leak",
                    "description": "敏感数据泄露",
                    "severity": "high",
                    "evidence": f"Sensitive data pattern '{pattern}' found in response"
                })
        
        return threats
    
    def detect_denial_of_service(self, request: Dict[str, Any], system_metrics: Dict[str, Any]) -> List[Dict[str, Any]]:
        """检测拒绝服务威胁"""
        threats = []
        
        # 检查高频调用
        if "call_frequency" in system_metrics and system_metrics["call_frequency"] > 1000:
            threats.append({
                "type": "denial_of_service",
                "name": "high_frequency_calls",
                "description": "高频调用",
                "severity": "medium",
                "evidence": f"Call frequency {system_metrics['call_frequency']} exceeds threshold 1000"
            })
        
        # 检查资源使用
        if "cpu_usage" in system_metrics and system_metrics["cpu_usage"] > 90:
            threats.append({
                "type": "denial_of_service",
                "name": "resource_exhaustion",
                "description": "资源耗尽",
                "severity": "high",
                "evidence": f"CPU usage {system_metrics['cpu_usage']}% exceeds threshold 90%"
            })
        
        return threats
    
    def detect_elevation_of_privilege(self, request: Dict[str, Any], user_permissions: List[str]) -> List[Dict[str, Any]]:
        """检测权限提升威胁"""
        threats = []
        
        # 检查未授权访问
        requested_tool = request.get("tool_name")
        if requested_tool and requested_tool not in user_permissions:
            threats.append({
                "type": "elevation_of_privilege",
                "name": "unauthorized_access",
                "description": "未授权访问",
                "severity": "high",
                "evidence": f"User {request.get('user_id')} attempted to access unauthorized tool {requested_tool}"
            })
        
        return threats
    
    def detect_threats(self, request: Dict[str, Any], response: Dict[str, Any] = None, system_metrics: Dict[str, Any] = None, user_permissions: List[str] = None) -> List[Dict[str, Any]]:
        """检测所有类型的威胁"""
        threats = []
        
        # 检测身份伪造
        threats.extend(self.detect_spoofing(request))
        
        # 检测篡改
        threats.extend(self.detect_tampering(request))
        
        # 检测信息泄露
        if response:
            threats.extend(self.detect_information_disclosure(response))
        
        # 检测拒绝服务
        if system_metrics:
            threats.extend(self.detect_denial_of_service(request, system_metrics))
        
        # 检测权限提升
        if user_permissions:
            threats.extend(self.detect_elevation_of_privilege(request, user_permissions))
        
        return threats
    
    def verify_signature(self, request: Dict[str, Any]) -> bool:
        """验证数字签名"""
        # 简化实现，实际需要使用真实的签名验证算法
        signature = request.get("signature")
        if not signature:
            return False
        
        # 计算请求体的哈希值
        request_body = {k: v for k, v in request.items() if k != "signature"}
        request_str = json.dumps(request_body, sort_keys=True)
        calculated_hash = hashlib.sha256(request_str.encode()).hexdigest()
        
        # 验证签名是否与哈希值匹配
        return signature == calculated_hash
    
    def calculate_parameters_hash(self, parameters: Dict[str, Any]) -> str:
        """计算参数的哈希值"""
        params_str = json.dumps(parameters, sort_keys=True)
        return hashlib.sha256(params_str.encode()).hexdigest()
    
    def get_valid_clients(self) -> List[str]:
        """获取有效的客户端列表"""
        # 简化实现，实际需要从配置或数据库中获取
        return ["client-001", "client-002", "client-003"]

# 测试代码
async def test_threat_detector():
    detector = MCPThreatDetector()
    
    # 正常请求
    normal_request = {
        "request_id": "req-001",
        "client_id": "client-001",
        "tool_name": "echo",
        "parameters": {"message": "hello world"},
        "parameters_hash": "a591a6d40bf420404a011733cfb7b190d62c65bf0bcda32b57b277d9ad9f146e",
        "signature": "a591a6d40bf420404a011733cfb7b190d62c65bf0bcda32b57b277d9ad9f146e"
    }
    
    # 异常请求（伪造身份）
    spoofing_request = {
        "request_id": "req-002",
        "client_id": "unknown-client",
        "tool_name": "echo",
        "parameters": {"message": "hello world"}
    }
    
    # 异常请求（篡改参数）
    tampering_request = {
        "request_id": "req-003",
        "client_id": "client-001",
        "tool_name": "echo",
        "parameters": {"message": "modified message"},
        "parameters_hash": "wrong-hash",
        "signature": "wrong-signature"
    }
    
    # 测试正常请求
    normal_threats = detector.detect_threats(normal_request, user_permissions=["echo"])
    print(f"Normal request threats: {normal_threats}")
    
    # 测试伪造身份请求
    spoofing_threats = detector.detect_threats(spoofing_request, user_permissions=["echo"])
    print(f"Spoofing request threats: {spoofing_threats}")
    
    # 测试篡改参数请求
    tampering_threats = detector.detect_threats(tampering_request, user_permissions=["echo"])
    print(f"Tampering request threats: {tampering_threats}")
    
    # 测试信息泄露响应
    sensitive_response = {
        "response_id": "resp-001",
        "status": "success",
        "result": {"password": "secret123", "message": "operation successful"}
    }
    sensitive_threats = detector.detect_information_disclosure(sensitive_response)
    print(f"Sensitive response threats: {sensitive_threats}")

if __name__ == "__main__":
    import asyncio
    asyncio.run(test_threat_detector())

from typing import Dict, List, Any
from enum import Enum

class ProbabilityLevel(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

class ImpactLevel(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

class RiskLevel(Enum):
    LOW = "low"
    MEDIUM_LOW = "medium_low"
    MEDIUM_HIGH = "medium_high"
    HIGH = "high"

class PriorityLevel(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

class RiskAssessor:
    def __init__(self):
        # 初始化风险评估矩阵
        self.risk_matrix = {
            (ProbabilityLevel.LOW, ImpactLevel.LOW): RiskLevel.LOW,
            (ProbabilityLevel.LOW, ImpactLevel.MEDIUM): RiskLevel.MEDIUM_LOW,
            (ProbabilityLevel.LOW, ImpactLevel.HIGH): RiskLevel.MEDIUM_HIGH,
            (ProbabilityLevel.MEDIUM, ImpactLevel.LOW): RiskLevel.LOW,
            (ProbabilityLevel.MEDIUM, ImpactLevel.MEDIUM): RiskLevel.MEDIUM_LOW,
            (ProbabilityLevel.MEDIUM, ImpactLevel.HIGH): RiskLevel.MEDIUM_HIGH,
            (ProbabilityLevel.HIGH, ImpactLevel.LOW): RiskLevel.MEDIUM_LOW,
            (ProbabilityLevel.HIGH, ImpactLevel.MEDIUM): RiskLevel.MEDIUM_HIGH,
            (ProbabilityLevel.HIGH, ImpactLevel.HIGH): RiskLevel.HIGH
        }
        
        # 初始化优先级映射
        self.priority_map = {
            RiskLevel.LOW: PriorityLevel.LOW,
            RiskLevel.MEDIUM_LOW: PriorityLevel.MEDIUM,
            RiskLevel.MEDIUM_HIGH: PriorityLevel.MEDIUM,
            RiskLevel.HIGH: PriorityLevel.HIGH
        }
    
    def assess_risk(self, probability: ProbabilityLevel, impact: ImpactLevel) -> Dict[str, Any]:
        """评估风险等级和优先级"""
        risk_level = self.risk_matrix.get((probability, impact), RiskLevel.LOW)
        priority = self.priority_map.get(risk_level, PriorityLevel.LOW)
        
        return {
            "probability": probability.value,
            "impact": impact.value,
            "risk_level": risk_level.value,
            "priority": priority.value
        }
    
    def assess_threat_risk(self, threat: Dict[str, Any], historical_data: List[Dict[str, Any]] = None) -> Dict[str, Any]:
        """评估威胁的风险等级和优先级"""
        # 简化实现，实际需要基于历史数据和威胁特征计算概率和影响
        threat_type = threat.get("type", "")
        severity = threat.get("severity", "medium")
        
        # 基于威胁类型和严重程度，估计概率和影响
        if threat_type == "spoofing":
            probability = ProbabilityLevel.MEDIUM
            impact = ImpactLevel.HIGH if severity == "high" else ImpactLevel.MEDIUM
        elif threat_type == "tampering":
            probability = ProbabilityLevel.MEDIUM
            impact = ImpactLevel.MEDIUM
        elif threat_type == "information_disclosure":
            probability = ProbabilityLevel.LOW
            impact = ImpactLevel.HIGH if severity == "high" else ImpactLevel.MEDIUM
        elif threat_type == "denial_of_service":
            probability = ProbabilityLevel.HIGH
            impact = ImpactLevel.MEDIUM if severity == "medium" else ImpactLevel.HIGH
        elif threat_type == "elevation_of_privilege":
            probability = ProbabilityLevel.HIGH
            impact = ImpactLevel.HIGH
        else:
            probability = ProbabilityLevel.MEDIUM
            impact = ImpactLevel.MEDIUM
        
        # 如果有历史数据，可以基于历史发生频率调整概率
        if historical_data:
            threat_count = sum(1 for t in historical_data if t.get("type") == threat_type)
            total_threats = len(historical_data)
            if threat_count / total_threats > 0.5:
                probability = ProbabilityLevel.HIGH
            elif threat_count / total_threats < 0.1:
                probability = ProbabilityLevel.LOW
        
        return self.assess_risk(probability, impact)
    
    def prioritize_threats(self, threats: List[Dict[str, Any]], historical_data: List[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """对威胁进行优先级排序"""
        # 为每个威胁添加风险评估结果
        threats_with_risk = []
        for threat in threats:
            risk_assessment = self.assess_threat_risk(threat, historical_data)
            threat_with_risk = {
                **threat,
                "risk_assessment": risk_assessment
            }
            threats_with_risk.append(threat_with_risk)
        
        # 按优先级排序
        priority_order = {
            PriorityLevel.HIGH.value: 0,
            PriorityLevel.MEDIUM.value: 1,
            PriorityLevel.LOW.value: 2
        }
        
        return sorted(
            threats_with_risk,
            key=lambda x: priority_order[x["risk_assessment"]["priority"]]
        )

# 测试代码
async def test_risk_assessor():
    assessor = RiskAssessor()
    
    # 测试风险矩阵
    risk1 = assessor.assess_risk(ProbabilityLevel.HIGH, ImpactLevel.HIGH)
    print(f"High probability, High impact: {risk1}")
    
    risk2 = assessor.assess_risk(ProbabilityLevel.LOW, ImpactLevel.LOW)
    print(f"Low probability, Low impact: {risk2}")
    
    risk3 = assessor.assess_risk(ProbabilityLevel.MEDIUM, ImpactLevel.HIGH)
    print(f"Medium probability, High impact: {risk3}")
    
    # 测试威胁风险评估
    threat = {
        "type": "elevation_of_privilege",
        "name": "unauthorized_access",
        "description": "未授权访问",
        "severity": "high",
        "evidence": "User attempted to access unauthorized tool"
    }
    
    threat_risk = assessor.assess_threat_risk(threat)
    print(f"Threat risk assessment: {threat_risk}")
    
    # 测试威胁优先级排序
    threats = [
        {
            "type": "tampering",
            "name": "modified_parameters",
            "description": "修改的调用参数",
            "severity": "medium",
            "evidence": "Parameters hash mismatch"
        },
        {
            "type": "elevation_of_privilege",
            "name": "unauthorized_access",
            "description": "未授权访问",
            "severity": "high",
            "evidence": "User attempted to access unauthorized tool"
        },
        {
            "type": "denial_of_service",
            "name": "high_frequency_calls",
            "description": "高频调用",
            "severity": "medium",
            "evidence": "Call frequency exceeds threshold"
        }
    ]
    
    prioritized_threats = assessor.prioritize_threats(threats)
    print(f"Prioritized threats: {prioritized_threats}")

if __name__ == "__main__":
    import asyncio
    asyncio.run(test_risk_assessor())

from typing import Dict, List, Any
from enum import Enum

class MitigationLevel(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

class MitigationStrategy:
    def __init__(self):
        # 初始化缓解策略库
        self.strategies = {
            "spoofing": [
                {
                    "name": "implement_mfa",
                    "description": "实施多因素认证",
                    "level": MitigationLevel.HIGH,
                    "type": "technical",
                    "implementation_cost": "high"
                },
                {
                    "name": "use_strong_authentication",
                    "description": "使用强身份认证",
                    "level": MitigationLevel.MEDIUM,
                    "type": "technical",
                    "implementation_cost": "medium"
                },
                {
                    "name": "regularly_rotate_credentials",
                    "description": "定期轮换凭证",
                    "level": MitigationLevel.MEDIUM,
                    "type": "management",
                    "implementation_cost": "low"
                }
            ],
            "tampering": [
                {
                    "name": "use_digital_signatures",
                    "description": "使用数字签名验证数据完整性",
                    "level": MitigationLevel.HIGH,
                    "type": "technical",
                    "implementation_cost": "medium"
                },
                {
                    "name": "implement_secure_transmission",
                    "description": "实施安全传输",
                    "level": MitigationLevel.MEDIUM,
                    "type": "technical",
                    "implementation_cost": "medium"
                },
                {
                    "name": "regularly_monitor_data_integrity",
                    "description": "定期监控数据完整性",
                    "level": MitigationLevel.MEDIUM,
                    "type": "management",
                    "implementation_cost": "low"
                }
            ],
            "information_disclosure": [
                {
                    "name": "encrypt_sensitive_data",
                    "description": "加密敏感数据",
                    "level": MitigationLevel.HIGH,
                    "type": "technical",
                    "implementation_cost": "high"
                },
                {
                    "name": "implement_access_control",
                    "description": "实施访问控制",
                    "level": MitigationLevel.MEDIUM,
                    "type": "technical",
                    "implementation_cost": "medium"
                },
                {
                    "name": "implement_data_masking",
                    "description": "实施数据脱敏",
                    "level": MitigationLevel.MEDIUM,
                    "type": "technical",
                    "implementation_cost": "medium"
                }
            ],
            "denial_of_service": [
                {
                    "name": "implement_rate_limiting",
                    "description": "实施速率限制",
                    "level": MitigationLevel.HIGH,
                    "type": "technical",
                    "implementation_cost": "medium"
                },
                {
                    "name": "deploy_load_balancer",
                    "description": "部署负载均衡器",
                    "level": MitigationLevel.MEDIUM,
                    "type": "technical",
                    "implementation_cost": "high"
                },
                {
                    "name": "implement_auto_scaling",
                    "description": "实施自动扩展",
                    "level": MitigationLevel.MEDIUM,
                    "type": "technical",
                    "implementation_cost": "high"
                }
            ],
            "elevation_of_privilege": [
                {
                    "name": "implement_rbac",
                    "description": "实施基于角色的访问控制",
                    "level": MitigationLevel.HIGH,
                    "type": "technical",
                    "implementation_cost": "medium"
                },
                {
                    "name": "enforce_least_privilege",
                    "description": "实施最小权限原则",
                    "level": MitigationLevel.HIGH,
                    "type": "management",
                    "implementation_cost": "low"
                },
                {
                    "name": "regularly_review_permissions",
                    "description": "定期审查权限",
                    "level": MitigationLevel.MEDIUM,
                    "type": "management",
                    "implementation_cost": "low"
                }
            ]
        }
    
    def get_mitigation_strategies(self, threat_type: str, risk_level: str = "medium") -> List[Dict[str, Any]]:
        """获取特定威胁类型的缓解策略"""
        strategies = self.strategies.get(threat_type, [])
        
        # 根据风险等级过滤策略
        if risk_level == "high":
            return [s for s in strategies if s["level"] in [MitigationLevel.HIGH, MitigationLevel.MEDIUM]]
        elif risk_level == "medium":
            return [s for s in strategies if s["level"] in [MitigationLevel.MEDIUM]]
        else:
            return [s for s in strategies if s["level"] == MitigationLevel.LOW]
    
    def generate_mitigation_plan(self, threats: List[Dict[str, Any]], prioritized: bool = True) -> List[Dict[str, Any]]:
        """生成缓解计划"""
        mitigation_plan = []
        
        # 为每个威胁生成缓解策略
        for threat in threats:
            threat_type = threat.get("type", "")
            risk_level = threat.get("risk_assessment", {}).get("risk_level", "medium")
            
            strategies = self.get_mitigation_strategies(threat_type, risk_level)
            
            for strategy in strategies:
                mitigation_plan.append({
                    "threat_id": threat.get("id", "unknown"),
                    "threat_type": threat_type,
                    "threat_description": threat.get("description", ""),
                    "strategy": strategy["name"],
                    "strategy_description": strategy["description"],
                    "strategy_type": strategy["type"],
                    "level": strategy["level"].value,
                    "implementation_cost": strategy["implementation_cost"],
                    "risk_level": risk_level,
                    "priority": threat.get("risk_assessment", {}).get("priority", "medium")
                })
        
        # 如果需要优先级排序，按优先级和风险等级排序
        if prioritized:
            priority_order = {"high": 0, "medium": 1, "low": 2}
            level_order = {"high": 0, "medium": 1, "low": 2}
            
            mitigation_plan = sorted(
                mitigation_plan,
                key=lambda x: (priority_order[x["priority"]], level_order[x["level"]])
            )
        
        return mitigation_plan
    
    def evaluate_strategy_effectiveness(self, strategy: Dict[str, Any], historical_data: List[Dict[str, Any]]) -> float:
        """评估缓解策略的有效性"""
        # 简化实现，实际需要基于历史数据计算策略的有效性
        strategy_name = strategy["name"]
        
        # 基于策略类型和实施成本，估计有效性
        if strategy["type"] == "technical":
            if strategy["implementation_cost"] == "high":
                return 0.9  # 高成本技术措施，有效性高
            elif strategy["implementation_cost"] == "medium":
                return 0.7  # 中等成本技术措施，有效性中等
            else:
                return 0.5  # 低成本技术措施，有效性低
        else:  # 管理或法律措施
            return 0.6  # 管理和法律措施，有效性中等

# 测试代码
async def test_mitigation_strategy():
    strategy_manager = MitigationStrategy()
    
    # 获取特定威胁类型的缓解策略
    spoofing_strategies = strategy_manager.get_mitigation_strategies("spoofing", "high")
    print(f"Spoofing mitigation strategies: {spoofing_strategies}")
    
    # 生成缓解计划
    threats = [
        {
            "id": "threat-001",
            "type": "elevation_of_privilege",
            "description": "未授权访问",
            "severity": "high",
            "risk_assessment": {"risk_level": "high", "priority": "high"}
        },
        {
            "id": "threat-002",
            "type": "information_disclosure",
            "description": "敏感数据泄露",
            "severity": "high",
            "risk_assessment": {"risk_level": "medium_high", "priority": "medium"}
        },
        {
            "id": "threat-003",
            "type": "tampering",
            "description": "修改调用参数",
            "severity": "medium",
            "risk_assessment": {"risk_level": "medium", "priority": "medium"}
        }
    ]
    
    mitigation_plan = strategy_manager.generate_mitigation_plan(threats)
    print(f"Mitigation plan: {mitigation_plan}")
    
    # 评估策略有效性
    strategy = {
        "name": "implement_mfa",
        "description": "实施多因素认证",
        "type": "technical",
        "implementation_cost": "high",
        "level": MitigationLevel.HIGH
    }
    effectiveness = strategy_manager.evaluate_strategy_effectiveness(strategy, [])
    print(f"Strategy effectiveness: {effectiveness}")

if __name__ == "__main__":
    import asyncio
    asyncio.run(test_mitigation_strategy())

detection:
  rules:
    spoofing:
      - name: invalid_signature
        description: 无效的数字签名
        severity: high
    tampering:
      - name: modified_parameters
        description: 修改的调用参数
        severity: medium
  thresholds:
    call_frequency: 1000
    cpu_usage: 90
    memory_usage: 90
  logging:
    level: INFO
    retention: 30d