""" MEMORY STORE: Persistent storage for adversarial learning data. Stores and retrieves historical attack results for pattern analysis. """ import os import json import sqlite3 import logging from typing import Dict, Any, List, Optional, Tuple, Union from datetime import datetime, timedelta from dataclasses import dataclass, asdict from pathlib import Path logger = logging.getLogger(__name__) @dataclass class AttackRecord: """Single attack execution record.""" attack_id: str attack_type: str attack_category: str target_model: str dataset: str prompt: str success: bool safety_score: float risk_level: str response_text: str response_length: int inference_time_ms: float timestamp: datetime metadata: Dict[str, Any] def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for storage.""" data = asdict(self) data['timestamp'] = self.timestamp.isoformat() return data @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'AttackRecord': """Create from dictionary.""" data['timestamp'] = datetime.fromisoformat(data['timestamp']) return cls(**data) @dataclass class PatternMetrics: """Metrics for attack patterns.""" attack_type: str total_attempts: int successful_attempts: int success_rate: float avg_safety_score: float avg_response_length: int avg_inference_time: float last_updated: datetime weak_indicators: List[str] strong_indicators: List[str] class MemoryStore: """ Persistent memory store for adversarial learning data. Stores historical attack results and provides analysis capabilities. """ def __init__(self, db_path: Optional[str] = None, enable_persistence: bool = True): """ Initialize memory store. Args: db_path: Path to SQLite database enable_persistence: Whether to persist data to disk """ self.enable_persistence = enable_persistence self.cache_dir = Path("cache/learning") if enable_persistence else None if db_path: self.db_path = Path(db_path) else: self.db_path = Path("data/learning_memory.db") if enable_persistence else ":memory:" # Create directories if self.enable_persistence and self.cache_dir: self.cache_dir.mkdir(parents=True, exist_ok=True) self.db_path.parent.mkdir(parents=True, exist_ok=True) # Initialize database self._init_database() # In-memory cache for frequent access self._attack_cache: List[AttackRecord] = [] self._pattern_cache: Dict[str, PatternMetrics] = {} logger.info(f"Memory store initialized: {self.db_path}") def _init_database(self): """Initialize SQLite database with required tables.""" try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() # Attacks table cursor.execute(""" CREATE TABLE IF NOT EXISTS attacks ( attack_id TEXT PRIMARY KEY, attack_type TEXT NOT NULL, attack_category TEXT NOT NULL, target_model TEXT NOT NULL, dataset TEXT NOT NULL, prompt TEXT NOT NULL, success BOOLEAN NOT NULL, safety_score REAL NOT NULL, risk_level TEXT NOT NULL, response_text TEXT, response_length INTEGER, inference_time_ms REAL, timestamp TEXT NOT NULL, metadata TEXT ) """) # Create indexes separately cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_timestamp ON attacks(timestamp)") cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_attack_type ON attacks(attack_type)") cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_target_model ON attacks(target_model)") cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_success ON attacks(success)") # Pattern metrics table cursor.execute(""" CREATE TABLE IF NOT EXISTS pattern_metrics ( attack_type TEXT PRIMARY KEY, total_attempts INTEGER NOT NULL, successful_attempts INTEGER NOT NULL, success_rate REAL NOT NULL, avg_safety_score REAL NOT NULL, avg_response_length INTEGER NOT NULL, avg_inference_time REAL NOT NULL, last_updated TEXT NOT NULL, weak_indicators TEXT, strong_indicators TEXT ) """) # Learning insights table cursor.execute(""" CREATE TABLE IF NOT EXISTS learning_insights ( insight_id TEXT PRIMARY KEY, insight_type TEXT NOT NULL, content TEXT NOT NULL, confidence REAL NOT NULL, created_at TEXT NOT NULL, applies_to_attack_types TEXT ) """) # Create indexes for learning insights cursor.execute("CREATE INDEX IF NOT EXISTS idx_learning_insights_type ON learning_insights(insight_type)") cursor.execute("CREATE INDEX IF NOT EXISTS idx_learning_insights_created_at ON learning_insights(created_at)") conn.commit() logger.info("Database initialized successfully") except Exception as e: logger.error(f"Failed to initialize database: {e}") raise def store_attack(self, attack_record: AttackRecord) -> bool: """ Store a single attack record. Args: attack_record: Attack record to store Returns: True if stored successfully """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() data = attack_record.to_dict() cursor.execute(""" INSERT OR REPLACE INTO attacks (attack_id, attack_type, attack_category, target_model, dataset, prompt, success, safety_score, risk_level, response_text, response_length, inference_time_ms, timestamp, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( data['attack_id'], data['attack_type'], data['attack_category'], data['target_model'], data['dataset'], data['prompt'], data['success'], data['safety_score'], data['risk_level'], data['response_text'], data['response_length'], data['inference_time_ms'], data['timestamp'], json.dumps(data['metadata']) )) conn.commit() # Update cache self._attack_cache.append(attack_record) logger.debug(f"Stored attack record: {attack_record.attack_id}") return True except Exception as e: logger.error(f"Failed to store attack record: {e}") return False def store_batch_attacks(self, attack_records: List[AttackRecord]) -> int: """ Store multiple attack records efficiently. Args: attack_records: List of attack records to store Returns: Number of records stored successfully """ if not attack_records: return 0 try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() stored_count = 0 for record in attack_records: data = record.to_dict() cursor.execute(""" INSERT OR REPLACE INTO attacks (attack_id, attack_type, attack_category, target_model, dataset, prompt, success, safety_score, risk_level, response_text, response_length, inference_time_ms, timestamp, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( data['attack_id'], data['attack_type'], data['attack_category'], data['target_model'], data['dataset'], data['prompt'], data['success'], data['safety_score'], data['risk_level'], data['response_text'], data['response_length'], data['inference_time_ms'], data['timestamp'], json.dumps(data['metadata']) )) stored_count += 1 conn.commit() # Update cache self._attack_cache.extend(attack_records) logger.info(f"Stored {stored_count} attack records") return stored_count except Exception as e: logger.error(f"Failed to store batch attacks: {e}") return 0 def get_attacks_by_model(self, model_name: str, limit: Optional[int] = None) -> List[AttackRecord]: """ Retrieve attacks for a specific model. Args: model_name: Target model name limit: Maximum number of records to retrieve Returns: List of attack records """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() query = "SELECT * FROM attacks WHERE target_model = ? ORDER BY timestamp DESC" params = [model_name] if limit: query += " LIMIT ?" params.append(limit) cursor.execute(query, params) rows = cursor.fetchall() # Convert to AttackRecord objects attacks = [] for row in rows: # Convert row to dict columns = [desc[0] for desc in cursor.description] data = dict(zip(columns, row)) data['metadata'] = json.loads(data['metadata']) if data['metadata'] else {} attacks.append(AttackRecord.from_dict(data)) return attacks except Exception as e: logger.error(f"Failed to retrieve attacks for model {model_name}: {e}") return [] def get_attacks_by_type(self, attack_type: str, limit: Optional[int] = None) -> List[AttackRecord]: """ Retrieve attacks of a specific type. Args: attack_type: Attack type to retrieve limit: Maximum number of records to retrieve Returns: List of attack records """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() query = "SELECT * FROM attacks WHERE attack_type = ? ORDER BY timestamp DESC" params = [attack_type] if limit: query += " LIMIT ?" params.append(limit) cursor.execute(query, params) rows = cursor.fetchall() # Convert to AttackRecord objects attacks = [] for row in rows: columns = [desc[0] for desc in cursor.description] data = dict(zip(columns, row)) data['metadata'] = json.loads(data['metadata']) if data['metadata'] else {} attacks.append(AttackRecord.from_dict(data)) return attacks except Exception as e: logger.error(f"Failed to retrieve attacks for type {attack_type}: {e}") return [] def get_recent_attacks(self, hours: int = 24, limit: Optional[int] = None) -> List[AttackRecord]: """ Retrieve recent attacks within specified time window. Args: hours: Number of hours to look back limit: Maximum number of records to retrieve Returns: List of recent attack records """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cutoff_time = (datetime.now() - timedelta(hours=hours)).isoformat() query = "SELECT * FROM attacks WHERE timestamp >= ? ORDER BY timestamp DESC" params = [cutoff_time] if limit: query += " LIMIT ?" params.append(limit) cursor.execute(query, params) rows = cursor.fetchall() # Convert to AttackRecord objects attacks = [] for row in rows: columns = [desc[0] for desc in cursor.description] data = dict(zip(columns, row)) data['metadata'] = json.loads(data['metadata']) if data['metadata'] else {} attacks.append(AttackRecord.from_dict(data)) return attacks except Exception as e: logger.error(f"Failed to retrieve recent attacks: {e}") return [] def get_attack_statistics(self, model_name: Optional[str] = None) -> Dict[str, Any]: """ Get comprehensive attack statistics. Args: model_name: Optional model filter Returns: Dictionary with attack statistics """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() # Base query where_clause = "WHERE target_model = ?" if model_name else "" params = [model_name] if model_name else [] # Total attacks cursor.execute(f"SELECT COUNT(*) FROM attacks {where_clause}", params) total_attacks = cursor.fetchone()[0] # Success rate cursor.execute(f""" SELECT SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) as successes, AVG(safety_score) as avg_safety, AVG(response_length) as avg_length, AVG(inference_time_ms) as avg_time FROM attacks {where_clause} """, params) success_data = cursor.fetchone() # Attack type breakdown cursor.execute(f""" SELECT attack_type, COUNT(*) as count, SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) as successes FROM attacks {where_clause} GROUP BY attack_type ORDER BY count DESC """, params) type_breakdown = cursor.fetchall() # Recent activity recent_cutoff = (datetime.now() - timedelta(hours=24)).isoformat() cursor.execute(f""" SELECT COUNT(*) FROM attacks WHERE timestamp >= ? {('AND target_model = ?' if model_name else '')} """, [recent_cutoff] + ([model_name] if model_name else [])) recent_count = cursor.fetchone()[0] return { "total_attacks": total_attacks, "successful_attacks": success_data[0] if success_data else 0, "success_rate": (success_data[0] / total_attacks * 100) if total_attacks > 0 else 0, "avg_safety_score": success_data[1] or 0, "avg_response_length": success_data[2] or 0, "avg_inference_time_ms": success_data[3] or 0, "attack_type_breakdown": [ { "type": row[0], "count": row[1], "successes": row[2], "success_rate": (row[2] / row[1] * 100) if row[1] > 0 else 0 } for row in type_breakdown ], "recent_attacks_24h": recent_count, "model_filter": model_name } except Exception as e: logger.error(f"Failed to get attack statistics: {e}") return {} def update_pattern_metrics(self, metrics: PatternMetrics) -> bool: """ Update pattern metrics for an attack type. Args: metrics: Pattern metrics to store Returns: True if updated successfully """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO pattern_metrics (attack_type, total_attempts, successful_attempts, success_rate, avg_safety_score, avg_response_length, avg_inference_time, last_updated, weak_indicators, strong_indicators) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( metrics.attack_type, metrics.total_attempts, metrics.successful_attempts, metrics.success_rate, metrics.avg_safety_score, metrics.avg_response_length, metrics.avg_inference_time, metrics.last_updated.isoformat(), json.dumps(metrics.weak_indicators), json.dumps(metrics.strong_indicators) )) conn.commit() # Update cache self._pattern_cache[metrics.attack_type] = metrics logger.debug(f"Updated pattern metrics for {metrics.attack_type}") return True except Exception as e: logger.error(f"Failed to update pattern metrics: {e}") return False def get_pattern_metrics(self, attack_type: Optional[str] = None) -> Union[PatternMetrics, Dict[str, PatternMetrics]]: """ Retrieve pattern metrics. Args: attack_type: Specific attack type, or None for all Returns: PatternMetrics object or dictionary of all metrics """ try: with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() if attack_type: cursor.execute("SELECT * FROM pattern_metrics WHERE attack_type = ?", [attack_type]) row = cursor.fetchone() if row: columns = [desc[0] for desc in cursor.description] data = dict(zip(columns, row)) return PatternMetrics( attack_type=data['attack_type'], total_attempts=data['total_attempts'], successful_attempts=data['successful_attempts'], success_rate=data['success_rate'], avg_safety_score=data['avg_safety_score'], avg_response_length=data['avg_response_length'], avg_inference_time=data['avg_inference_time'], last_updated=datetime.fromisoformat(data['last_updated']), weak_indicators=json.loads(data['weak_indicators']) if data['weak_indicators'] else [], strong_indicators=json.loads(data['strong_indicators']) if data['strong_indicators'] else [] ) else: return None else: cursor.execute("SELECT * FROM pattern_metrics ORDER BY success_rate DESC") rows = cursor.fetchall() metrics_dict = {} for row in rows: columns = [desc[0] for desc in cursor.description] data = dict(zip(columns, row)) metrics_dict[data['attack_type']] = PatternMetrics( attack_type=data['attack_type'], total_attempts=data['total_attempts'], successful_attempts=data['successful_attempts'], success_rate=data['success_rate'], avg_safety_score=data['avg_safety_score'], avg_response_length=data['avg_response_length'], avg_inference_time=data['avg_inference_time'], last_updated=datetime.fromisoformat(data['last_updated']), weak_indicators=json.loads(data['weak_indicators']) if data['weak_indicators'] else [], strong_indicators=json.loads(data['strong_indicators']) if data['strong_indicators'] else [] ) return metrics_dict except Exception as e: logger.error(f"Failed to retrieve pattern metrics: {e}") return {} if not attack_type else None def cleanup_old_data(self, days_to_keep: int = 30) -> int: """ Clean up old attack data to manage storage. Args: days_to_keep: Number of days to keep data (0 = no cleanup) Returns: Number of records removed """ try: # Check if cleanup is disabled if days_to_keep <= 0: logger.info("Data cleanup is disabled (LEARNING_CLEANUP_DAYS=0)") return 0 with sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cutoff_time = (datetime.now() - timedelta(days=days_to_keep)).isoformat() cursor.execute("DELETE FROM attacks WHERE timestamp < ?", [cutoff_time]) removed_count = cursor.rowcount conn.commit() # Clear cache self._attack_cache = [r for r in self._attack_cache if r.timestamp >= datetime.fromisoformat(cutoff_time)] logger.info(f"Cleaned up {removed_count} old attack records") return removed_count except Exception as e: logger.error(f"Failed to cleanup old data: {e}") return 0 def export_data(self, output_path: str, model_name: Optional[str] = None) -> bool: """ Export attack data to JSON file. Args: output_path: Output file path model_name: Optional model filter Returns: True if exported successfully """ try: attacks = self.get_attacks_by_model(model_name) if model_name else self.get_recent_attacks(hours=24*30) # 30 days export_data = { "export_timestamp": datetime.now().isoformat(), "model_filter": model_name, "total_records": len(attacks), "attacks": [attack.to_dict() for attack in attacks] } with open(output_path, 'w') as f: json.dump(export_data, f, indent=2) logger.info(f"Exported {len(attacks)} attack records to {output_path}") return True except Exception as e: logger.error(f"Failed to export data: {e}") return False # Singleton instance for global access _memory_store_instance: Optional[MemoryStore] = None def get_memory_store(db_path: Optional[str] = None, enable_persistence: bool = True, db_session=None): """ Get memory store instance. Args: db_path: Path to SQLite database (deprecated, use PostgreSQL instead) enable_persistence: Whether to enable persistence db_session: PostgreSQL database session (preferred) Returns: MemoryStore instance """ global _memory_store_instance # If database session is provided, reset global instance to ensure PostgreSQL usage if db_session is not None: _memory_store_instance = None if _memory_store_instance is None: # Use PostgreSQL if session is provided, otherwise fallback to SQLite if db_session is not None: from .postgresql_memory_store import PostgreSQLMemoryStore _memory_store_instance = PostgreSQLMemoryStore(db_session, enable_persistence) else: logger.warning("⚠️ No database session provided, using SQLite (not recommended for production)") _memory_store_instance = MemoryStore(db_path, enable_persistence) return _memory_store_instance