""" Federated Learning System for Cyber-LLM Enables secure collaborative learning across multiple organizations without sharing raw data. Author: Muzan Sano """ import asyncio import json import hashlib import logging from datetime import datetime, timedelta from typing import Dict, List, Optional, Tuple, Any, Union from dataclasses import dataclass, asdict from enum import Enum import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset from transformers import AutoTokenizer, AutoModelForCausalLM import websockets import ssl from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC import base64 import os from ..utils.logging_system import CyberLLMLogger from ..utils.secrets_manager import SecretsManager from .online_learning import LearningEvent, LearningEventType # Configure logging logger = CyberLLMLogger(__name__).get_logger() class FederatedRole(Enum): """Roles in federated learning network""" COORDINATOR = "coordinator" # Central coordination server PARTICIPANT = "participant" # Individual organization VALIDATOR = "validator" # Validates model updates class FederatedMessageType(Enum): """Types of messages in federated learning protocol""" JOIN_REQUEST = "join_request" JOIN_RESPONSE = "join_response" MODEL_UPDATE = "model_update" AGGREGATION_REQUEST = "aggregation_request" AGGREGATION_RESPONSE = "aggregation_response" VALIDATION_REQUEST = "validation_request" VALIDATION_RESPONSE = "validation_response" HEARTBEAT = "heartbeat" @dataclass class FederatedParticipant: """Information about a federated learning participant""" participant_id: str organization: str public_key: str last_seen: datetime contribution_weight: float = 1.0 # Weight based on data quality/quantity trust_score: float = 1.0 # Trust level (0-1) specialization: List[str] = None # Areas of expertise def __post_init__(self): if self.specialization is None: self.specialization = [] @dataclass class FederatedMessage: """Structure for federated learning messages""" message_id: str sender_id: str recipient_id: str # "broadcast" for all participants message_type: FederatedMessageType payload: Dict[str, Any] timestamp: datetime signature: Optional[str] = None encrypted: bool = False def to_dict(self) -> Dict[str, Any]: data = asdict(self) data['timestamp'] = self.timestamp.isoformat() data['message_type'] = self.message_type.value return data class SecureCommunicationManager: """Manages secure communication between federated participants""" def __init__(self, participant_id: str): self.participant_id = participant_id self.encryption_key = None self.participants_keys: Dict[str, str] = {} def generate_encryption_key(self, password: bytes) -> None: """Generate encryption key from password""" salt = os.urandom(16) kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, ) key = base64.urlsafe_b64encode(kdf.derive(password)) self.encryption_key = Fernet(key) def encrypt_message(self, message: Dict[str, Any]) -> bytes: """Encrypt message payload""" if self.encryption_key is None: raise ValueError("Encryption key not set") message_bytes = json.dumps(message).encode() return self.encryption_key.encrypt(message_bytes) def decrypt_message(self, encrypted_data: bytes) -> Dict[str, Any]: """Decrypt message payload""" if self.encryption_key is None: raise ValueError("Encryption key not set") decrypted_bytes = self.encryption_key.decrypt(encrypted_data) return json.loads(decrypted_bytes.decode()) def sign_message(self, message: Dict[str, Any]) -> str: """Create digital signature for message""" message_str = json.dumps(message, sort_keys=True) return hashlib.sha256(message_str.encode()).hexdigest() def verify_signature(self, message: Dict[str, Any], signature: str) -> bool: """Verify message digital signature""" expected_signature = self.sign_message(message) return expected_signature == signature class ModelAggregator: """Handles secure model aggregation in federated learning""" def __init__(self, aggregation_method: str = "fedavg"): self.aggregation_method = aggregation_method self.model_updates: List[Dict[str, torch.Tensor]] = [] self.participant_weights: List[float] = [] def add_model_update(self, model_state: Dict[str, torch.Tensor], weight: float = 1.0): """Add a model update from a participant""" self.model_updates.append(model_state) self.participant_weights.append(weight) def aggregate_models(self) -> Dict[str, torch.Tensor]: """Aggregate multiple model updates using specified method""" if not self.model_updates: raise ValueError("No model updates to aggregate") if self.aggregation_method == "fedavg": return self._federated_averaging() elif self.aggregation_method == "weighted_avg": return self._weighted_averaging() else: raise ValueError(f"Unknown aggregation method: {self.aggregation_method}") def _federated_averaging(self) -> Dict[str, torch.Tensor]: """Standard federated averaging aggregation""" if not self.model_updates: return {} # Get parameter names from first model param_names = self.model_updates[0].keys() aggregated_params = {} total_weight = sum(self.participant_weights) for param_name in param_names: weighted_sum = torch.zeros_like(self.model_updates[0][param_name]) for i, model_update in enumerate(self.model_updates): weight = self.participant_weights[i] / total_weight weighted_sum += weight * model_update[param_name] aggregated_params[param_name] = weighted_sum return aggregated_params def _weighted_averaging(self) -> Dict[str, torch.Tensor]: """Weighted averaging based on participant trust scores""" # Similar to federated averaging but uses trust scores return self._federated_averaging() def clear_updates(self): """Clear accumulated model updates""" self.model_updates.clear() self.participant_weights.clear() class FederatedLearningCoordinator: """Coordinates federated learning across multiple participants""" def __init__(self, coordinator_id: str, port: int = 8765, min_participants: int = 3, aggregation_rounds: int = 10): self.coordinator_id = coordinator_id self.port = port self.min_participants = min_participants self.aggregation_rounds = aggregation_rounds # Participant management self.participants: Dict[str, FederatedParticipant] = {} self.connected_clients = set() # Learning state self.current_round = 0 self.model_aggregator = ModelAggregator() self.global_model = None self.round_results: List[Dict[str, Any]] = [] # Communication self.comm_manager = SecureCommunicationManager(coordinator_id) self.server = None logger.info(f"FederatedLearningCoordinator initialized: {coordinator_id}") async def start_coordinator(self): """Start the federated learning coordinator server""" try: # Create SSL context for secure communication ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) # In production, load proper certificates self.server = await websockets.serve( self.handle_client, "localhost", self.port, ssl=None # Enable SSL in production ) logger.info(f"Federated learning coordinator started on port {self.port}") # Start coordination loop await self.coordination_loop() except Exception as e: logger.error(f"Failed to start coordinator: {str(e)}") async def handle_client(self, websocket, path): """Handle incoming client connections""" try: self.connected_clients.add(websocket) logger.info("New participant connected") async for message in websocket: await self.process_message(websocket, message) except websockets.exceptions.ConnectionClosed: logger.info("Participant disconnected") finally: self.connected_clients.discard(websocket) async def process_message(self, websocket, raw_message: str): """Process incoming message from participant""" try: message_data = json.loads(raw_message) message = FederatedMessage(**message_data) if message.message_type == FederatedMessageType.JOIN_REQUEST: await self.handle_join_request(websocket, message) elif message.message_type == FederatedMessageType.MODEL_UPDATE: await self.handle_model_update(websocket, message) elif message.message_type == FederatedMessageType.HEARTBEAT: await self.handle_heartbeat(websocket, message) else: logger.warning(f"Unknown message type: {message.message_type}") except Exception as e: logger.error(f"Error processing message: {str(e)}") async def handle_join_request(self, websocket, message: FederatedMessage): """Handle participant join request""" try: participant_info = message.payload participant = FederatedParticipant( participant_id=message.sender_id, organization=participant_info.get('organization', 'unknown'), public_key=participant_info.get('public_key', ''), last_seen=datetime.now(), specialization=participant_info.get('specialization', []) ) self.participants[message.sender_id] = participant # Send join response response = FederatedMessage( message_id=f"join_resp_{datetime.now().timestamp()}", sender_id=self.coordinator_id, recipient_id=message.sender_id, message_type=FederatedMessageType.JOIN_RESPONSE, payload={ 'accepted': True, 'participant_id': message.sender_id, 'current_round': self.current_round }, timestamp=datetime.now() ) await websocket.send(json.dumps(response.to_dict())) logger.info(f"Participant {message.sender_id} joined from {participant.organization}") except Exception as e: logger.error(f"Error handling join request: {str(e)}") async def handle_model_update(self, websocket, message: FederatedMessage): """Handle model update from participant""" try: update_data = message.payload # Verify update integrity if not self.verify_model_update(update_data): logger.warning(f"Invalid model update from {message.sender_id}") return # Extract model parameters (in practice, this would be more complex) model_params = update_data.get('model_parameters', {}) participant_weight = self.participants[message.sender_id].contribution_weight # Add to aggregator self.model_aggregator.add_model_update(model_params, participant_weight) logger.info(f"Received model update from {message.sender_id}") # Check if ready for aggregation if len(self.model_aggregator.model_updates) >= self.min_participants: await self.perform_aggregation() except Exception as e: logger.error(f"Error handling model update: {str(e)}") async def handle_heartbeat(self, websocket, message: FederatedMessage): """Handle heartbeat from participant""" if message.sender_id in self.participants: self.participants[message.sender_id].last_seen = datetime.now() def verify_model_update(self, update_data: Dict[str, Any]) -> bool: """Verify the integrity and validity of a model update""" # Implement security checks: # 1. Digital signature verification # 2. Parameter bounds checking # 3. Differential privacy validation # 4. Anomaly detection required_fields = ['model_parameters', 'training_metrics', 'data_size'] return all(field in update_data for field in required_fields) async def perform_aggregation(self): """Perform model aggregation and distribute updated model""" try: logger.info(f"Starting aggregation round {self.current_round}") # Aggregate model updates aggregated_params = self.model_aggregator.aggregate_models() # Update global model (simplified) self.global_model = aggregated_params # Broadcast updated model to all participants await self.broadcast_updated_model(aggregated_params) # Record round results round_result = { 'round': self.current_round, 'participants': len(self.model_aggregator.model_updates), 'timestamp': datetime.now().isoformat(), 'aggregation_method': self.model_aggregator.aggregation_method } self.round_results.append(round_result) # Clean up for next round self.model_aggregator.clear_updates() self.current_round += 1 logger.info(f"Aggregation round {self.current_round - 1} completed") except Exception as e: logger.error(f"Error performing aggregation: {str(e)}") async def broadcast_updated_model(self, model_params: Dict[str, Any]): """Broadcast updated global model to all participants""" message = FederatedMessage( message_id=f"agg_resp_{datetime.now().timestamp()}", sender_id=self.coordinator_id, recipient_id="broadcast", message_type=FederatedMessageType.AGGREGATION_RESPONSE, payload={ 'global_model_parameters': model_params, 'round': self.current_round, 'participants_count': len(self.participants) }, timestamp=datetime.now() ) # Send to all connected clients if self.connected_clients: message_str = json.dumps(message.to_dict()) await asyncio.gather( *[client.send(message_str) for client in self.connected_clients], return_exceptions=True ) async def coordination_loop(self): """Main coordination loop""" while True: try: # Check participant health await self.check_participant_health() # Trigger periodic aggregation if needed await self.check_aggregation_trigger() await asyncio.sleep(30) # Check every 30 seconds except Exception as e: logger.error(f"Error in coordination loop: {str(e)}") await asyncio.sleep(60) # Wait longer on error async def check_participant_health(self): """Check health of all participants""" current_time = datetime.now() timeout_threshold = timedelta(minutes=5) inactive_participants = [] for participant_id, participant in self.participants.items(): if current_time - participant.last_seen > timeout_threshold: inactive_participants.append(participant_id) # Remove inactive participants for participant_id in inactive_participants: del self.participants[participant_id] logger.info(f"Removed inactive participant: {participant_id}") async def check_aggregation_trigger(self): """Check if aggregation should be triggered""" # Trigger based on time or number of updates if (len(self.model_aggregator.model_updates) >= self.min_participants and len(self.model_aggregator.model_updates) < len(self.participants)): # Wait for more participants or trigger after timeout pass def get_federation_status(self) -> Dict[str, Any]: """Get current federation status""" return { 'coordinator_id': self.coordinator_id, 'current_round': self.current_round, 'active_participants': len(self.participants), 'connected_clients': len(self.connected_clients), 'pending_updates': len(self.model_aggregator.model_updates), 'total_rounds': len(self.round_results), 'participants': { pid: { 'organization': p.organization, 'last_seen': p.last_seen.isoformat(), 'trust_score': p.trust_score, 'specialization': p.specialization } for pid, p in self.participants.items() } } class FederatedLearningParticipant: """Federated learning participant (individual organization)""" def __init__(self, participant_id: str, organization: str, coordinator_url: str = "ws://localhost:8765", specialization: List[str] = None): self.participant_id = participant_id self.organization = organization self.coordinator_url = coordinator_url self.specialization = specialization or [] # Local model and data self.local_model = None self.local_data: List[LearningEvent] = [] # Communication self.comm_manager = SecureCommunicationManager(participant_id) self.websocket = None self.connected = False logger.info(f"FederatedLearningParticipant initialized: {participant_id}") async def join_federation(self) -> bool: """Join the federated learning federation""" try: self.websocket = await websockets.connect(self.coordinator_url) self.connected = True # Send join request join_message = FederatedMessage( message_id=f"join_{datetime.now().timestamp()}", sender_id=self.participant_id, recipient_id="coordinator", message_type=FederatedMessageType.JOIN_REQUEST, payload={ 'organization': self.organization, 'public_key': 'participant_public_key', # Replace with actual key 'specialization': self.specialization }, timestamp=datetime.now() ) await self.websocket.send(json.dumps(join_message.to_dict())) # Start message handling loop asyncio.create_task(self.message_handler()) logger.info(f"Joined federation as {self.participant_id}") return True except Exception as e: logger.error(f"Failed to join federation: {str(e)}") return False async def message_handler(self): """Handle incoming messages from coordinator""" try: async for message in self.websocket: await self.process_coordinator_message(message) except websockets.exceptions.ConnectionClosed: self.connected = False logger.info("Connection to coordinator lost") async def process_coordinator_message(self, raw_message: str): """Process message from coordinator""" try: message_data = json.loads(raw_message) message = FederatedMessage(**message_data) if message.message_type == FederatedMessageType.AGGREGATION_RESPONSE: await self.handle_global_model_update(message) elif message.message_type == FederatedMessageType.JOIN_RESPONSE: await self.handle_join_response(message) else: logger.info(f"Received message type: {message.message_type}") except Exception as e: logger.error(f"Error processing coordinator message: {str(e)}") async def handle_global_model_update(self, message: FederatedMessage): """Handle updated global model from coordinator""" try: global_params = message.payload.get('global_model_parameters', {}) round_number = message.payload.get('round', 0) # Update local model with global parameters if self.local_model and global_params: # In practice, this would update the actual model logger.info(f"Updated local model with global parameters from round {round_number}") # Optionally trigger new local training round await self.train_local_model() except Exception as e: logger.error(f"Error handling global model update: {str(e)}") async def handle_join_response(self, message: FederatedMessage): """Handle join response from coordinator""" payload = message.payload if payload.get('accepted', False): logger.info("Successfully joined federation") else: logger.error("Federation join request rejected") async def train_local_model(self): """Train local model and send update to coordinator""" if not self.local_data: logger.warning("No local data available for training") return try: # Simulate local training (implement actual training logic) logger.info(f"Training local model with {len(self.local_data)} samples") # Generate model update (simplified) model_update = { 'model_parameters': {}, # Actual model parameters 'training_metrics': { 'loss': 0.1, 'accuracy': 0.9, 'samples': len(self.local_data) }, 'data_size': len(self.local_data) } # Send update to coordinator await self.send_model_update(model_update) except Exception as e: logger.error(f"Error training local model: {str(e)}") async def send_model_update(self, model_update: Dict[str, Any]): """Send model update to coordinator""" if not self.connected: logger.error("Not connected to coordinator") return try: update_message = FederatedMessage( message_id=f"update_{datetime.now().timestamp()}", sender_id=self.participant_id, recipient_id="coordinator", message_type=FederatedMessageType.MODEL_UPDATE, payload=model_update, timestamp=datetime.now() ) await self.websocket.send(json.dumps(update_message.to_dict())) logger.info("Sent model update to coordinator") except Exception as e: logger.error(f"Error sending model update: {str(e)}") def add_local_data(self, learning_events: List[LearningEvent]): """Add learning events to local training data""" self.local_data.extend(learning_events) logger.info(f"Added {len(learning_events)} learning events to local data") # Factory functions def create_federated_coordinator(**kwargs) -> FederatedLearningCoordinator: """Create federated learning coordinator""" return FederatedLearningCoordinator(**kwargs) def create_federated_participant(**kwargs) -> FederatedLearningParticipant: """Create federated learning participant""" return FederatedLearningParticipant(**kwargs)