MCP 操作的全链路审计

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

class AuditBlock:
    def __init__(self, block_id: str, previous_hash: str, 
                 audit_data: List[Dict], timestamp: float, 
                 nonce: int = 0):
        self.block_id = block_id
        self.previous_hash = previous_hash
        self.audit_data = audit_data
        self.timestamp = timestamp
        self.nonce = nonce
        self.hash = self.calculate_hash()
    
    def calculate_hash(self) -> str:
        """计算区块哈希值"""
        block_data = {
            "block_id": self.block_id,
            "previous_hash": self.previous_hash,
            "audit_data": self.audit_data,
            "timestamp": self.timestamp,
            "nonce": self.nonce
        }
        
        block_json = json.dumps(block_data, sort_keys=True)
        return hashlib.sha256(block_json.encode()).hexdigest()
    
    def mine_block(self, difficulty: int) -> bool:
        """挖矿生成新区块"""
        target = "0" * difficulty
        while self.hash[:difficulty] != target:
            self.nonce += 1
            self.hash = self.calculate_hash()
        
        logging.info(f"Block mined: {self.block_id} with nonce: {self.nonce}")
        return True
    
    def to_dict(self) -> Dict:
        """转换为字典"""
        return {
            "block_id": self.block_id,
            "previous_hash": self.previous_hash,
            "audit_data": self.audit_data,
            "timestamp": self.timestamp,
            "nonce": self.nonce,
            "hash": self.hash
        }

class BlockchainAuditStorage:
    def __init__(self, difficulty: int = 4):
        self.difficulty = difficulty
        self.chain = [self._create_genesis_block()]
        self.pending_audit_data = []
        self.logger = logging.getLogger("mcp_blockchain_audit")
    
    def _create_genesis_block(self) -> AuditBlock:
        """创建创世区块"""
        genesis_block = AuditBlock(
            block_id="genesis",
            previous_hash="0" * 64,
            audit_data=[],
            timestamp=datetime.now().timestamp()
        )
        return genesis_block
    
    def get_latest_block(self) -> AuditBlock:
        """获取最新区块"""
        return self.chain[-1]
    
    def add_audit_data(self, audit_data: Dict):
        """添加审计数据到待处理列表"""
        self.pending_audit_data.append(audit_data)
        self.logger.info(f"Added audit data: {audit_data.get('audit_id')}")
    
    def mine_pending_audit_data(self, miner_id: str) -> AuditBlock:
        """挖矿处理待处理的审计数据"""
        if not self.pending_audit_data:
            return None
        
        latest_block = self.get_latest_block()
        new_block = AuditBlock(
            block_id=f"block_{len(self.chain)}",
            previous_hash=latest_block.hash,
            audit_data=self.pending_audit_data.copy(),
            timestamp=datetime.now().timestamp()
        )
        
        # 挖矿生成新区块
        new_block.mine_block(self.difficulty)
        
        # 添加新区块到区块链
        self.chain.append(new_block)
        
        # 清空待处理审计数据
        self.pending_audit_data = []
        
        self.logger.info(f"Mined new block: {new_block.block_id} with {len(new_block.audit_data)} audit records")
        return new_block
    
    def is_chain_valid(self) -> bool:
        """验证区块链的有效性"""
        for i in range(1, len(self.chain)):
            current_block = self.chain[i]
            previous_block = self.chain[i-1]
            
            # 验证当前区块哈希
            if current_block.hash != current_block.calculate_hash():
                self.logger.error(f"Block {current_block.block_id} has invalid hash")
                return False
            
            # 验证前一个区块哈希引用
            if current_block.previous_hash != previous_block.hash:
                self.logger.error(f"Block {current_block.block_id} has invalid previous hash")
                return False
        
        return True
    
    def query_audit_data(self, query: Dict) -> List[Dict]:
        """查询审计数据"""
        results = []
        
        for block in self.chain:
            for audit_data in block.audit_data:
                match = True
                for key, value in query.items():
                    if key not in audit_data or audit_data[key] != value:
                        match = False
                        break
                if match:
                    results.append(audit_data)
        
        return results
    
    def get_chain_info(self) -> Dict:
        """获取区块链信息"""
        return {
            "chain_length": len(self.chain),
            "difficulty": self.difficulty,
            "pending_audit_data_count": len(self.pending_audit_data),
            "latest_block": self.get_latest_block().to_dict()
        }

# mcp_ai_audit.py
from typing import List, Dict, Optional
import logging
import numpy as np
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
import joblib
from datetime import datetime

class AuditFeatureExtractor:
    def __init__(self):
        self.logger = logging.getLogger("mcp_ai_audit.feature")
    
    def extract_features(self, audit_data: Dict) -> np.ndarray:
        """从审计数据中提取特征"""
        # 提取基础特征
        features = []
        
        # 操作类型编码
        action_type_map = {
            "tool_call": 0,
            "auth": 1,
            "config_change": 2,
            "status_check": 3
        }
        action_type = action_type_map.get(audit_data.get("action_type", "unknown"), 4)
        features.append(action_type)
        
        # 调用频率（这里简化处理，实际应基于历史数据计算）
        call_frequency = audit_data.get("call_frequency", 0)
        features.append(call_frequency)
        
        # 参数数量
        parameters = audit_data.get("parameters", {})
        param_count = len(parameters)
        features.append(param_count)
        
        # 执行时间
        duration = audit_data.get("duration_ms", 0)
        features.append(duration)
        
        # 用户类型编码
        user_type_map = {
            "admin": 0,
            "user": 1,
            "service": 2
        }
        user_type = user_type_map.get(audit_data.get("user_type", "unknown"), 3)
        features.append(user_type)
        
        return np.array(features).reshape(1, -1)
    
    def extract_batch_features(self, audit_data_list: List[Dict]) -> np.ndarray:
        """批量提取特征"""
        features_list = []
        for audit_data in audit_data_list:
            features = self.extract_features(audit_data)
            features_list.append(features)
        return np.vstack(features_list)

class AIAuditAnalyzer:
    def __init__(self, model_path: Optional[str] = None):
        self.feature_extractor = AuditFeatureExtractor()
        
        # 创建或加载模型
        if model_path:
            self.model = joblib.load(model_path)
        else:
            self.model = Pipeline([
                ("scaler", StandardScaler()),
                ("clf", IsolationForest(n_estimators=100, contamination=0.01, random_state=42))
            ])
        
        self.logger = logging.getLogger("mcp_ai_audit.analyzer")
    
    def train(self, training_data: List[Dict]):
        """训练模型"""
        self.logger.info(f"Training model with {len(training_data)} samples")
        
        # 提取特征
        X = self.feature_extractor.extract_batch_features(training_data)
        
        # 训练模型
        self.model.fit(X)
        self.logger.info("Model trained successfully")
    
    def save_model(self, model_path: str):
        """保存模型"""
        joblib.dump(self.model, model_path)
        self.logger.info(f"Model saved to {model_path}")
    
    def detect_anomaly(self, audit_data: Dict) -> Dict:
        """检测异常操作"""
        # 提取特征
        features = self.feature_extractor.extract_features(audit_data)
        
        # 预测异常
        prediction = self.model.predict(features)[0]
        
        # 计算异常分数
        anomaly_score = self.model.score_samples(features)[0]
        
        result = {
            "audit_id": audit_data.get("audit_id"),
            "is_anomaly": prediction == -1,
            "anomaly_score": anomaly_score,
            "timestamp": datetime.now().timestamp()
        }
        
        if prediction == -1:
            self.logger.warning(f"Anomaly detected: {audit_data.get('audit_id')}, score: {anomaly_score}")
        
        return result
    
    def detect_batch_anomalies(self, audit_data_list: List[Dict]) -> List[Dict]:
        """批量检测异常操作"""
        results = []
        for audit_data in audit_data_list:
            result = self.detect_anomaly(audit_data)
            results.append(result)
        return results
    
    def update_model(self, new_data: List[Dict]):
        """更新模型"""
        # 实现模型增量更新逻辑
        self.logger.info(f"Updating model with {len(new_data)} new samples")
        
        # 提取特征
        X_new = self.feature_extractor.extract_batch_features(new_data)
        
        # 增量更新模型
        # 这里简化处理，实际应使用支持增量学习的模型
        self.model.fit(X_new)
        
        self.logger.info("Model updated successfully")

# mcp_realtime_audit_alert.py
from typing import List, Dict, Optional
import logging
from datetime import datetime
import asyncio

class AlertRule:
    def __init__(self, rule_id: str, name: str, description: str, 
                 condition: Dict, severity: str, action: str):
        self.rule_id = rule_id
        self.name = name
        self.description = description
        self.condition = condition
        self.severity = severity  # low, medium, high, critical
        self.action = action  # email, sms, webhook, all
    
    def to_dict(self) -> Dict:
        return {
            "rule_id": self.rule_id,
            "name": self.name,
            "description": self.description,
            "condition": self.condition,
            "severity": self.severity,
            "action": self.action
        }
    
    def evaluate(self, audit_data: Dict) -> bool:
        """评估告警规则"""
        for key, value in self.condition.items():
            if key not in audit_data:
                return False
            
            if isinstance(value, dict):
                # 处理比较操作
                if "operator" in value and "value" in value:
                    operator = value["operator"]
                    compare_value = value["value"]
                    actual_value = audit_data[key]
                    
                    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
            else:
                # 直接比较
                if audit_data[key] != value:
                    return False
        
        return True

class Alert:
    def __init__(self, alert_id: str, rule_id: str, audit_id: str, 
                 severity: str, message: str, timestamp: float, 
                 audit_data: Dict):
        self.alert_id = alert_id
        self.rule_id = rule_id
        self.audit_id = audit_id
        self.severity = severity
        self.message = message
        self.timestamp = timestamp
        self.audit_data = audit_data
    
    def to_dict(self) -> Dict:
        return {
            "alert_id": self.alert_id,
            "rule_id": self.rule_id,
            "audit_id": self.audit_id,
            "severity": self.severity,
            "message": self.message,
            "timestamp": self.timestamp,
            "audit_data": self.audit_data
        }

class RealTimeAuditAlertSystem:
    def __init__(self):
        self.rules = {}
        self.alerts = []
        self.logger = logging.getLogger("mcp_realtime_audit_alert")
    
    def add_rule(self, rule: AlertRule):
        """添加告警规则"""
        self.rules[rule.rule_id] = rule
        self.logger.info(f"Added alert rule: {rule.rule_id}")
    
    def remove_rule(self, rule_id: str):
        """移除告警规则"""
        if rule_id in self.rules:
            del self.rules[rule_id]
            self.logger.info(f"Removed alert rule: {rule_id}")
    
    def update_rule(self, rule: AlertRule):
        """更新告警规则"""
        self.rules[rule.rule_id] = rule
        self.logger.info(f"Updated alert rule: {rule.rule_id}")
    
    def evaluate_audit_data(self, audit_data: Dict) -> List[Alert]:
        """评估审计数据，生成告警"""
        triggered_alerts = []
        
        for rule in self.rules.values():
            if rule.evaluate(audit_data):
                # 生成告警
                alert = Alert(
                    alert_id=f"alert_{audit_data.get('audit_id')}_{rule.rule_id}_{datetime.now().timestamp()}",
                    rule_id=rule.rule_id,
                    audit_id=audit_data.get("audit_id"),
                    severity=rule.severity,
                    message=f"Rule {rule.name} triggered for audit {audit_data.get('audit_id')}",
                    timestamp=datetime.now().timestamp(),
                    audit_data=audit_data
                )
                
                triggered_alerts.append(alert)
                self.alerts.append(alert)
                
                # 发送告警通知
                self._send_alert(alert)
        
        return triggered_alerts
    
    def _send_alert(self, alert: Alert):
        """发送告警通知"""
        self.logger.warning(f"Sending alert: {alert.alert_id}, severity: {alert.severity}")
        
        # 根据告警规则的 action 字段，发送不同类型的通知
        if alert.action in ["email", "all"]:
            self._send_email_alert(alert)
        if alert.action in ["sms", "all"]:
            self._send_sms_alert(alert)
        if alert.action in ["webhook", "all"]:
            self._send_webhook_alert(alert)
    
    def _send_email_alert(self, alert: Alert):
        """发送邮件告警"""
        # 实现邮件告警发送逻辑
        self.logger.info(f"Sent email alert: {alert.alert_id}")
    
    def _send_sms_alert(self, alert: Alert):
        """发送短信告警"""
        # 实现短信告警发送逻辑
        self.logger.info(f"Sent SMS alert: {alert.alert_id}")
    
    def _send_webhook_alert(self, alert: Alert):
        """发送 Webhook 告警"""
        # 实现 Webhook 告警发送逻辑
        self.logger.info(f"Sent webhook alert: {alert.alert_id}")
    
    def get_alerts(self, filters: Optional[Dict] = None) -> List[Alert]:
        """获取告警列表"""
        if not filters:
            return self.alerts
        
        filtered_alerts = []
        for alert in self.alerts:
            match = True
            for key, value in filters.items():
                if key not in alert.__dict__ or alert.__dict__[key] != value:
                    match = False
                    break
            if match:
                filtered_alerts.append(alert)
        
        return filtered_alerts
    
    def get_alert_stats(self) -> Dict:
        """获取告警统计信息"""
        stats = {
            "total_alerts": len(self.alerts),
            "severity_counts": {
                "low": 0,
                "medium": 0,
                "high": 0,
                "critical": 0
            },
            "rule_trigger_counts": {}
        }
        
        for alert in self.alerts:
            stats["severity_counts"][alert.severity] += 1
            
            if alert.rule_id not in stats["rule_trigger_counts"]:
                stats["rule_trigger_counts"][alert.rule_id] = 0
            stats["rule_trigger_counts"][alert.rule_id] += 1
        
        return stats

# mcp_audit_collector.py
from typing import Dict, Optional
import logging
import json
from datetime import datetime
import uuid

class AuditCollector:
    def __init__(self, component_type: str, component_id: str):
        self.component_type = component_type  # client, server, host
        self.component_id = component_id
        self.logger = logging.getLogger(f"mcp_audit_collector.{component_type}")
    
    def generate_audit_log(self, action_type: str, context: Dict, 
                         result: str, status_code: int, 
                         message: str, duration_ms: int = 0) -> Dict:
        """生成审计日志"""
        audit_id = str(uuid.uuid4())
        
        audit_log = {
            "audit_id": audit_id,
            "component_type": self.component_type,
            "component_id": self.component_id,
            "action_type": action_type,
            "timestamp": datetime.now().timestamp(),
            "context": context,
            "result": result,
            "status_code": status_code,
            "message": message,
            "duration_ms": duration_ms
        }
        
        # 添加额外的上下文信息
        audit_log["call_frequency"] = self._calculate_call_frequency(action_type)
        audit_log["user_type"] = context.get("user_type", "unknown")
        
        self.logger.info(f"Generated audit log: {audit_id}")
        return audit_log
    
    def _calculate_call_frequency(self, action_type: str) -> float:
        """计算调用频率"""
        # 实现调用频率计算逻辑
        # 这里简化处理，返回默认值
        return 1.0
    
    def send_audit_log(self, audit_log: Dict, audit_bus_url: str):
        """发送审计日志到审计数据总线"""
        # 实现审计日志发送逻辑
        # 这里简化处理，直接打印日志
        self.logger.info(f"Sending audit log to {audit_bus_url}: {audit_log.get('audit_id')}")
    
    def generate_and_send(self, action_type: str, context: Dict, 
                        result: str, status_code: int, 
                        message: str, duration_ms: int = 0, 
                        audit_bus_url: str = "http://localhost:8080/audit"):
        """生成并发送审计日志"""
        audit_log = self.generate_audit_log(action_type, context, result, 
                                          status_code, message, duration_ms)
        self.send_audit_log(audit_log, audit_bus_url)
        return audit_log

# 示例使用
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建审计采集器
    server_collector = AuditCollector(component_type="server", component_id="server-001")
    
    # 生成并发送审计日志
    context = {
        "client_id": "client-001",
        "tool_id": "file_reader",
        "parameters": {
            "file_path": "/data/test.txt",
            "mode": "read"
        },
        "user_id": "user-001",
        "user_type": "user"
    }
    
    server_collector.generate_and_send(
        action_type="tool_call",
        context=context,
        result="success",
        status_code=200,
        message="工具调用成功",
        duration_ms=100
    )

# 安装依赖
pip install kafka-python scikit-learn joblib numpy

# 配置审计系统
cp config.example.yaml config.yaml
# 编辑配置文件
vim config.yaml

# 启动审计数据总线（Kafka）
docker-compose up -d kafka zookeeper

# 启动审计数据处理器
python mcp_audit_processor.py --config config.yaml

# 启动区块链审计存储
python mcp_blockchain_audit.py --config config.yaml

# 启动 AI 审计分析引擎
python mcp_ai_audit.py --config config.yaml

# 启动实时审计告警系统
python mcp_realtime_audit_alert.py --config config.yaml

# 启动审计数据可视化系统（Kibana）
docker-compose up -d kibana