""" ATTACK OPTIMIZER: Optimizes attack generation based on learning insights. Modifies attack strategies to maximize success rates based on historical data. """ import logging from typing import Dict, Any, List, Optional, Tuple from datetime import datetime import random from dataclasses import dataclass from .learning_engine import LearningEngine, AttackStrategy, get_learning_engine from .memory_store import AttackRecord, get_memory_store logger = logging.getLogger(__name__) @dataclass class OptimizedAttackConfig: """Configuration for optimized attack generation.""" attack_type: str priority: float # 0.0 to 1.0 weight_multiplier: float # How much to weight this attack type prompt_modifications: List[str] target_models: List[str] success_prediction: float confidence: float reasoning: str last_updated: datetime class AttackOptimizer: """ Optimizes attack generation based on learning insights and historical patterns. Modifies attack strategies to maximize success rates. """ def __init__(self, learning_engine: Optional[LearningEngine] = None, db_session=None): """ Initialize attack optimizer. Args: learning_engine: Learning engine instance for accessing insights db_session: Database session for PostgreSQL memory store """ self.learning_engine = learning_engine or get_learning_engine(db_session=db_session) self.memory_store = get_memory_store(db_session=db_session) # Optimization parameters self.min_confidence_threshold = 0.6 # Minimum confidence to use insights self.success_rate_weight = 0.7 # Weight for historical success rate self.insight_weight = 0.3 # Weight for learning insights # Cache for optimized configurations self._config_cache: Dict[str, List[OptimizedAttackConfig]] = {} self._cache_expiry_minutes = 60 # Cache expires after 1 hour logger.info("Attack optimizer initialized") def optimize_attack_strategy(self, target_model: str, available_attack_types: List[str]) -> List[OptimizedAttackConfig]: """ Optimize attack strategy for a specific model. Args: target_model: Target model name available_attack_types: List of available attack types Returns: List of optimized attack configurations """ try: cache_key = f"{target_model}_{hash(tuple(available_attack_types))}" # Check cache if self._is_cache_valid(cache_key): logger.debug(f"Using cached optimization for {target_model}") return self._config_cache[cache_key] # Get adaptive strategy strategy = self.learning_engine.build_adaptive_strategy(target_model) if not strategy: logger.warning(f"No adaptive strategy available for {target_model}") return self._create_default_configurations(available_attack_types) # Get learning insights insights = self.learning_engine.generate_learning_insights(target_model) # Optimize based on strategy and insights optimized_configs = self._create_optimized_configurations(strategy, insights, available_attack_types) # Cache results self._config_cache[cache_key] = optimized_configs logger.info(f"Generated {len(optimized_configs)} optimized configurations for {target_model}") return optimized_configs except Exception as e: logger.error(f"Failed to optimize attack strategy for {target_model}: {e}") return self._create_default_configurations(available_attack_types) def get_attack_priorities(self, target_model: str) -> Dict[str, float]: """ Get priority weights for different attack types against a model. Args: target_model: Target model name Returns: Dictionary mapping attack types to priority weights (0.0 to 1.0) """ try: # Get adaptive strategy strategy = self.learning_engine.build_adaptive_strategy(target_model) if not strategy: return {} priorities = {} # High priority for primary attack types for attack_type in strategy.primary_attack_types: priorities[attack_type] = 0.9 # Medium priority for secondary attack types for attack_type in strategy.secondary_attack_types: priorities[attack_type] = 0.6 # Low priority for avoided attack types for attack_type in strategy.avoided_attack_types: priorities[attack_type] = 0.1 return priorities except Exception as e: logger.error(f"Failed to get attack priorities for {target_model}: {e}") return {} def optimize_prompt_for_attack(self, attack_type: str, base_prompt: str, target_model: str) -> str: """ Optimize a prompt for a specific attack type and target model. Args: attack_type: Type of attack base_prompt: Base prompt to optimize target_model: Target model name Returns: Optimized prompt """ try: # Get strategy for target model strategy = self.learning_engine.build_adaptive_strategy(target_model) if not strategy: return base_prompt # Get historical data for this attack type and model attacks = self.memory_store.get_attacks_by_model(target_model) relevant_attacks = [a for a in attacks if a.attack_type == attack_type and a.success] if not relevant_attacks: return base_prompt # Extract optimization patterns optimized_prompt = self._apply_prompt_optimizations(base_prompt, relevant_attacks, strategy) return optimized_prompt except Exception as e: logger.error(f"Failed to optimize prompt: {e}") return base_prompt def update_optimization_from_results(self, target_model: str, attack_results: List[AttackRecord]) -> bool: """ Update optimization based on new attack results. Args: target_model: Target model name attack_results: New attack results Returns: True if optimization updated successfully """ try: # Store results and update learning strategy_id = None if target_model in self.learning_engine._strategy_cache: strategy_id = self.learning_engine._strategy_cache[target_model].strategy_id if strategy_id: success = self.learning_engine.update_strategy_from_results(strategy_id, attack_results) # Clear cache to force re-optimization for cache_key in list(self._config_cache.keys()): if cache_key.startswith(target_model): del self._config_cache[cache_key] logger.info(f"Updated optimization for {target_model} from {len(attack_results)} results") return success else: logger.warning(f"No strategy found for {target_model}") return False except Exception as e: logger.error(f"Failed to update optimization: {e}") return False def get_optimization_report(self, target_model: str) -> Dict[str, Any]: """ Get comprehensive optimization report for a model. Args: target_model: Target model name Returns: Optimization report """ try: # Get strategy strategy = self.learning_engine.build_adaptive_strategy(target_model) # Get priorities priorities = self.get_attack_priorities(target_model) # Get insights insights = self.learning_engine.generate_learning_insights(target_model) # Get historical performance recent_attacks = self.memory_store.get_attacks_by_model(target_model, limit=100) report = { "target_model": target_model, "timestamp": datetime.now().isoformat(), "has_strategy": strategy is not None, "strategy_confidence": strategy.confidence if strategy else 0.0, "predicted_success_rate": strategy.success_prediction if strategy else 0.0, "attack_priorities": priorities, "primary_attack_types": strategy.primary_attack_types if strategy else [], "secondary_attack_types": strategy.secondary_attack_types if strategy else [], "avoided_attack_types": strategy.avoided_attack_types if strategy else [], "supporting_insights": len(insights), "recent_performance": { "total_attacks": len(recent_attacks), "success_rate": sum(1 for a in recent_attacks if a.success) / len(recent_attacks) * 100 if recent_attacks else 0, "avg_safety_score": sum(a.safety_score for a in recent_attacks) / len(recent_attacks) if recent_attacks else 0 }, "optimization_status": "active" if strategy else "inactive" } return report except Exception as e: logger.error(f"Failed to generate optimization report: {e}") return {"error": str(e)} def _create_optimized_configurations(self, strategy: AttackStrategy, insights: List, available_attack_types: List[str]) -> List[OptimizedAttackConfig]: """Create optimized attack configurations based on strategy and insights.""" configs = [] # Primary attack types (highest priority) for attack_type in strategy.primary_attack_types: if attack_type in available_attack_types: config = self._create_attack_config(attack_type, strategy, insights, priority=0.9) configs.append(config) # Secondary attack types (medium priority) for attack_type in strategy.secondary_attack_types: if attack_type in available_attack_types and attack_type not in [c.attack_type for c in configs]: config = self._create_attack_config(attack_type, strategy, insights, priority=0.6) configs.append(config) # Available attack types not in strategy (low priority) for attack_type in available_attack_types: if attack_type not in [c.attack_type for c in configs] and attack_type not in strategy.avoided_attack_types: config = self._create_attack_config(attack_type, strategy, insights, priority=0.3) configs.append(config) # Sort by priority configs.sort(key=lambda x: x.priority, reverse=True) return configs def _create_attack_config(self, attack_type: str, strategy: AttackStrategy, insights: List, priority: float) -> OptimizedAttackConfig: """Create optimized configuration for a specific attack type.""" # Get historical performance attacks = self.memory_store.get_attacks_by_model(strategy.target_model) relevant_attacks = [a for a in attacks if a.attack_type == attack_type] success_rate = 0.0 if relevant_attacks: success_rate = sum(1 for a in relevant_attacks if a.success) / len(relevant_attacks) * 100 # Calculate confidence based on data volume and strategy confidence data_confidence = min(1.0, len(relevant_attacks) / 20.0) # More data = higher confidence final_confidence = (data_confidence * 0.6) + (strategy.confidence * 0.4) # Generate prompt modifications prompt_modifications = self._generate_prompt_modifications(attack_type, relevant_attacks, insights) # Calculate weight multiplier weight_multiplier = self._calculate_weight_multiplier(attack_type, strategy, success_rate) # Generate reasoning reasoning = self._generate_config_reasoning(attack_type, strategy, success_rate, insights) return OptimizedAttackConfig( attack_type=attack_type, priority=priority, weight_multiplier=weight_multiplier, prompt_modifications=prompt_modifications, target_models=[strategy.target_model], success_prediction=success_rate, confidence=final_confidence, reasoning=reasoning, last_updated=datetime.now() ) def _generate_prompt_modifications(self, attack_type: str, relevant_attacks: List[AttackRecord], insights: List) -> List[str]: """Generate prompt modifications based on historical success patterns.""" modifications = [] if not relevant_attacks: return modifications # Analyze successful prompts successful_attacks = [a for a in relevant_attacks if a.success] if successful_attacks: # Length patterns avg_length = sum(len(a.prompt) for a in successful_attacks) / len(successful_attacks) if avg_length > 200: modifications.append("Use longer, more detailed prompts") elif avg_length < 100: modifications.append("Use shorter, more direct prompts") # Common keywords all_text = " ".join([a.prompt.lower() for a in successful_attacks]) words = all_text.split() word_freq = {} for word in words: word_freq[word] = word_freq.get(word, 0) + 1 # Find frequently used words common_words = [word for word, freq in sorted(word_freq.items(), key=lambda x: x[1], reverse=True)[:5]] if common_words: modifications.append(f"Consider using keywords: {', '.join(common_words[:3])}") # Add insight-based modifications for insight in insights: if attack_type in insight.applies_to_attack_types and insight.confidence >= 0.7: if "weakness" in insight.insight_type: modifications.append(f"Exploit identified weakness: {insight.content}") elif "pattern" in insight.insight_type: modifications.append(f"Apply successful pattern: {insight.content}") return modifications def _calculate_weight_multiplier(self, attack_type: str, strategy: AttackStrategy, success_rate: float) -> float: """Calculate weight multiplier for attack type.""" base_multiplier = 1.0 # Adjust based on strategy priority if attack_type in strategy.primary_attack_types: base_multiplier *= 1.5 elif attack_type in strategy.secondary_attack_types: base_multiplier *= 1.2 elif attack_type in strategy.avoided_attack_types: base_multiplier *= 0.3 # Adjust based on historical success rate if success_rate > 70: base_multiplier *= 1.3 elif success_rate < 30: base_multiplier *= 0.7 # Adjust based on strategy confidence if strategy.confidence > 0.8: base_multiplier *= 1.1 elif strategy.confidence < 0.5: base_multiplier *= 0.9 # Ensure reasonable bounds return max(0.1, min(3.0, base_multiplier)) def _generate_config_reasoning(self, attack_type: str, strategy: AttackStrategy, success_rate: float, insights: List) -> str: """Generate reasoning for attack configuration.""" reasoning_parts = [] # Base reasoning if attack_type in strategy.primary_attack_types: reasoning_parts.append(f"Primary attack type with predicted success rate of {success_rate:.1f}%") elif attack_type in strategy.secondary_attack_types: reasoning_parts.append(f"Secondary attack type with predicted success rate of {success_rate:.1f}%") else: reasoning_parts.append(f"Alternative attack type with predicted success rate of {success_rate:.1f}%") # Strategy-based reasoning if strategy.confidence > 0.8: reasoning_parts.append(f"High confidence strategy ({strategy.confidence:.1f})") # Insight-based reasoning relevant_insights = [i for i in insights if attack_type in i.applies_to_attack_types] if relevant_insights: reasoning_parts.append(f"Supported by {len(relevant_insights)} learning insights") # Data-based reasoning reasoning_parts.append(f"Based on {len(self.memory_store.get_attacks_by_type(attack_type))} historical attacks") return "; ".join(reasoning_parts) def _create_default_configurations(self, available_attack_types: List[str]) -> List[OptimizedAttackConfig]: """Create default configurations when no learning data is available.""" configs = [] for attack_type in available_attack_types: config = OptimizedAttackConfig( attack_type=attack_type, priority=0.5, # Default priority weight_multiplier=1.0, # Default weight prompt_modifications=["Default configuration - no learning data available"], target_models=[], # No specific targets success_prediction=50.0, # Default prediction confidence=0.3, # Low confidence without data reasoning="Default configuration - insufficient historical data", last_updated=datetime.now() ) configs.append(config) return configs def _apply_prompt_optimizations(self, base_prompt: str, successful_attacks: List[AttackRecord], strategy: AttackStrategy) -> str: """Apply optimizations to a prompt based on successful historical attacks.""" if not successful_attacks: return base_prompt optimized_prompt = base_prompt # Length optimization avg_length = sum(len(a.prompt) for a in successful_attacks) / len(successful_attacks) current_length = len(base_prompt) if current_length < avg_length * 0.8: # Prompt is too short, add detail optimized_prompt += " Please provide a detailed and comprehensive response." elif current_length > avg_length * 1.2: # Prompt is too long, make it more concise optimized_prompt = base_prompt[:int(avg_length)] + "..." # Keyword optimization all_successful_text = " ".join([a.prompt.lower() for a in successful_attacks]) words = all_successful_text.split() word_freq = {} for word in words: word_freq[word] = word_freq.get(word, 0) + 1 # Add common successful words if not present common_words = [word for word, freq in sorted(word_freq.items(), key=lambda x: x[1], reverse=True)[:3]] for word in common_words: if word not in optimized_prompt.lower(): optimized_prompt += f" Consider {word}." return optimized_prompt def _is_cache_valid(self, cache_key: str) -> bool: """Check if cached configuration is still valid.""" if cache_key not in self._config_cache: return False # Check if any configuration is older than cache expiry for config in self._config_cache[cache_key]: age_minutes = (datetime.now() - config.last_updated).total_seconds() / 60 if age_minutes > self._cache_expiry_minutes: return False return True # Singleton instance for global access _attack_optimizer_instance: Optional[AttackOptimizer] = None def get_attack_optimizer(learning_engine: Optional[LearningEngine] = None, db_session=None) -> AttackOptimizer: """ Get global attack optimizer instance. Args: learning_engine: Learning engine instance db_session: Database session for PostgreSQL memory store Returns: AttackOptimizer instance """ global _attack_optimizer_instance # If database session is provided, reset global instance to ensure PostgreSQL usage if db_session is not None: _attack_optimizer_instance = None if _attack_optimizer_instance is None: _attack_optimizer_instance = AttackOptimizer(learning_engine, db_session) return _attack_optimizer_instance