aunk/ransomshield-pe-malware

1"""
2Database Module - Experiment History Storage
3
4This module provides SQLite database functionality to store, retrieve,
5and compare experiment results from both supervised and unsupervised runs.
6"""
7
8import sqlite3
9import json
10import os
11from datetime import datetime
12from typing import List, Optional, Dict, Any
13
14
15# Default database path (same directory as this file)
16DEFAULT_DB_PATH = os.path.join(os.path.dirname(__file__), "experiments.db")
17
18
19class ExperimentDatabase:
20    """
21    SQLite database manager for storing experiment results.
22    
23    Stores all metrics, parameters, and metadata for each experiment run.
24    """
25    
26    def __init__(self, db_path: str = DEFAULT_DB_PATH):
27        """
28        Initialize the database connection.
29        
30        Args:
31            db_path: Path to the SQLite database file
32        """
33        self.db_path = db_path
34        self._init_database()
35    
36    def _init_database(self):
37        """Create the experiments table if it doesn't exist."""
38        conn = sqlite3.connect(self.db_path)
39        cursor = conn.cursor()
40        
41        cursor.execute('''
42            CREATE TABLE IF NOT EXISTS experiments (
43                id INTEGER PRIMARY KEY AUTOINCREMENT,
44                timestamp TEXT NOT NULL,
45                experiment_type TEXT NOT NULL,
46                features_used TEXT NOT NULL,
47                num_features INTEGER,
48                
49                -- Metrics
50                accuracy REAL,
51                precision REAL,
52                recall REAL,
53                specificity REAL,
54                f1_score REAL,
55                auc_roc REAL,
56                
57                -- Confusion Matrix
58                true_positives INTEGER,
59                true_negatives INTEGER,
60                false_positives INTEGER,
61                false_negatives INTEGER,
62                
63                -- Parameters (stored as JSON)
64                parameters TEXT,
65                
66                -- Optional notes
67                notes TEXT
68            )
69        ''')
70        
71        conn.commit()
72        conn.close()
73    
74    def save_experiment(self, results: Dict[str, Any], notes: str = "") -> int:
75        """
76        Save an experiment result to the database.
77        
78        Args:
79            results: Dictionary containing experiment results
80            notes: Optional notes about the experiment
81        
82        Returns:
83            The ID of the inserted experiment
84        """
85        conn = sqlite3.connect(self.db_path)
86        cursor = conn.cursor()
87        
88        metrics = results.get('metrics', {})
89        cm = results.get('confusion_matrix', {})
90        params = results.get('parameters', {})
91        
92        cursor.execute('''
93            INSERT INTO experiments (
94                timestamp, experiment_type, features_used, num_features,
95                accuracy, precision, recall, specificity, f1_score, auc_roc,
96                true_positives, true_negatives, false_positives, false_negatives,
97                parameters, notes
98            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
99        ''', (
100            datetime.now().isoformat(),
101            results.get('experiment_type', 'Unknown'),
102            results.get('features_used', 'Unknown'),
103            results.get('num_features', 0),
104            metrics.get('accuracy', 0),
105            metrics.get('precision', 0),
106            metrics.get('recall', 0),
107            metrics.get('specificity', 0),
108            metrics.get('f1_score', 0),
109            metrics.get('auc_roc', 0),
110            cm.get('tp', 0),
111            cm.get('tn', 0),
112            cm.get('fp', 0),
113            cm.get('fn', 0),
114            json.dumps(params),
115            notes
116        ))
117        
118        experiment_id = cursor.lastrowid
119        conn.commit()
120        conn.close()
121        
122        return experiment_id
123    
124    def get_all_experiments(self) -> List[Dict[str, Any]]:
125        """
126        Retrieve all experiments from the database.
127        
128        Returns:
129            List of experiment dictionaries
130        """
131        conn = sqlite3.connect(self.db_path)
132        conn.row_factory = sqlite3.Row
133        cursor = conn.cursor()
134        
135        cursor.execute('''
136            SELECT * FROM experiments ORDER BY timestamp DESC
137        ''')
138        
139        rows = cursor.fetchall()
140        conn.close()
141        
142        return [dict(row) for row in rows]
143    
144    def get_experiment_by_id(self, experiment_id: int) -> Optional[Dict[str, Any]]:
145        """
146        Retrieve a specific experiment by ID.
147        
148        Args:
149            experiment_id: The ID of the experiment
150        
151        Returns:
152            Experiment dictionary or None if not found
153        """
154        conn = sqlite3.connect(self.db_path)
155        conn.row_factory = sqlite3.Row
156        cursor = conn.cursor()
157        
158        cursor.execute('SELECT * FROM experiments WHERE id = ?', (experiment_id,))
159        row = cursor.fetchone()
160        conn.close()
161        
162        return dict(row) if row else None
163    
164    def filter_experiments(
165        self,
166        experiment_type: Optional[str] = None,
167        features_used: Optional[str] = None,
168        min_accuracy: Optional[float] = None,
169        min_auc: Optional[float] = None,
170        limit: int = 100
171    ) -> List[Dict[str, Any]]:
172        """
173        Filter experiments based on criteria.
174        
175        Args:
176            experiment_type: Filter by experiment type (e.g., "Supervised", "Unsupervised")
177            features_used: Filter by features used (e.g., "All Features", "Chosen Features")
178            min_accuracy: Minimum accuracy threshold
179            min_auc: Minimum AUC-ROC threshold
180            limit: Maximum number of results
181        
182        Returns:
183            List of matching experiment dictionaries
184        """
185        conn = sqlite3.connect(self.db_path)
186        conn.row_factory = sqlite3.Row
187        cursor = conn.cursor()
188        
189        query = "SELECT * FROM experiments WHERE 1=1"
190        params = []
191        
192        if experiment_type:
193            query += " AND experiment_type LIKE ?"
194            params.append(f"%{experiment_type}%")
195        
196        if features_used:
197            query += " AND features_used LIKE ?"
198            params.append(f"%{features_used}%")
199        
200        if min_accuracy is not None:
201            query += " AND accuracy >= ?"
202            params.append(min_accuracy)
203        
204        if min_auc is not None:
205            query += " AND auc_roc >= ?"
206            params.append(min_auc)
207        
208        query += " ORDER BY timestamp DESC LIMIT ?"
209        params.append(limit)
210        
211        cursor.execute(query, params)
212        rows = cursor.fetchall()
213        conn.close()
214        
215        return [dict(row) for row in rows]
216    
217    def delete_experiment(self, experiment_id: int) -> bool:
218        """
219        Delete an experiment by ID.
220        
221        Args:
222            experiment_id: The ID of the experiment to delete
223        
224        Returns:
225            True if deleted, False if not found
226        """
227        conn = sqlite3.connect(self.db_path)
228        cursor = conn.cursor()
229        
230        cursor.execute('DELETE FROM experiments WHERE id = ?', (experiment_id,))
231        deleted = cursor.rowcount > 0
232        
233        conn.commit()
234        conn.close()
235        
236        return deleted
237    
238    def clear_all_experiments(self) -> int:
239        """
240        Delete all experiments from the database.
241        
242        Returns:
243            Number of experiments deleted
244        """
245        conn = sqlite3.connect(self.db_path)
246        cursor = conn.cursor()
247        
248        cursor.execute('DELETE FROM experiments')
249        deleted = cursor.rowcount
250        
251        conn.commit()
252        conn.close()
253        
254        return deleted
255    
256    def get_best_experiment(self, metric: str = "accuracy") -> Optional[Dict[str, Any]]:
257        """
258        Get the best experiment based on a specific metric.
259        
260        Args:
261            metric: The metric to sort by (accuracy, precision, recall, f1_score, auc_roc)
262        
263        Returns:
264            Best experiment dictionary or None if no experiments
265        """
266        valid_metrics = ["accuracy", "precision", "recall", "specificity", "f1_score", "auc_roc"]
267        if metric not in valid_metrics:
268            metric = "accuracy"
269        
270        conn = sqlite3.connect(self.db_path)
271        conn.row_factory = sqlite3.Row
272        cursor = conn.cursor()
273        
274        cursor.execute(f'SELECT * FROM experiments ORDER BY {metric} DESC LIMIT 1')
275        row = cursor.fetchone()
276        conn.close()
277        
278        return dict(row) if row else None
279    
280    def get_experiment_count(self) -> int:
281        """Return the total number of experiments in the database."""
282        conn = sqlite3.connect(self.db_path)
283        cursor = conn.cursor()
284        
285        cursor.execute('SELECT COUNT(*) FROM experiments')
286        count = cursor.fetchone()[0]
287        conn.close()
288        
289        return count
290    
291    def format_experiment_summary(self, exp: Dict[str, Any]) -> str:
292        """
293        Format an experiment as a readable summary string.
294        
295        Args:
296            exp: Experiment dictionary
297        
298        Returns:
299            Formatted string summary
300        """
301        timestamp = exp.get('timestamp', '')[:19].replace('T', ' ')
302        return (
303            f"[ID: {exp['id']}] {timestamp}\n"
304            f"  Type: {exp['experiment_type']}\n"
305            f"  Features: {exp['features_used']} ({exp['num_features']})\n"
306            f"  Accuracy: {exp['accuracy']*100:.2f}%  |  "
307            f"AUC: {exp['auc_roc']:.4f}  |  "
308            f"F1: {exp['f1_score']*100:.2f}%\n"
309        )
310
311
312# Singleton instance for easy access
313_db_instance = None
314
315def get_database() -> ExperimentDatabase:
316    """Get the singleton database instance."""
317    global _db_instance
318    if _db_instance is None:
319        _db_instance = ExperimentDatabase()
320    return _db_instance
321
322
323if __name__ == "__main__":
324    # Test the database
325    db = ExperimentDatabase()
326    print(f"Database initialized at: {db.db_path}")
327    print(f"Total experiments: {db.get_experiment_count()}")
328