Add alerts/recording, WiFi/TSCM updates, optimize waterfall

utils/database.py

@@ -88,19 +88,65 @@ def init_db() -> None:
             ON signal_history(mode, device_id, timestamp)
         ''')
 
-        # Device correlation table
-        conn.execute('''
-            CREATE TABLE IF NOT EXISTS device_correlations (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                wifi_mac TEXT,
-                bt_mac TEXT,
-                confidence REAL,
-                first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                metadata TEXT,
-                UNIQUE(wifi_mac, bt_mac)
-            )
-        ''')
+        # Device correlation table
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS device_correlations (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                wifi_mac TEXT,
+                bt_mac TEXT,
+                confidence REAL,
+                first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                metadata TEXT,
+                UNIQUE(wifi_mac, bt_mac)
+            )
+        ''')
+
+        # Alert rules
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS alert_rules (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT NOT NULL,
+                mode TEXT,
+                event_type TEXT,
+                match TEXT,
+                severity TEXT DEFAULT 'medium',
+                enabled BOOLEAN DEFAULT 1,
+                notify TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        ''')
+
+        # Alert events
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS alert_events (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                rule_id INTEGER,
+                mode TEXT,
+                event_type TEXT,
+                severity TEXT DEFAULT 'medium',
+                title TEXT,
+                message TEXT,
+                payload TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (rule_id) REFERENCES alert_rules(id) ON DELETE SET NULL
+            )
+        ''')
+
+        # Session recordings
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS recording_sessions (
+                id TEXT PRIMARY KEY,
+                mode TEXT NOT NULL,
+                label TEXT,
+                started_at TIMESTAMP NOT NULL,
+                stopped_at TIMESTAMP,
+                file_path TEXT NOT NULL,
+                event_count INTEGER DEFAULT 0,
+                size_bytes INTEGER DEFAULT 0,
+                metadata TEXT
+            )
+        ''')
 
         # Users table for authentication
         conn.execute('''
@@ -131,20 +177,29 @@ def init_db() -> None:
         # =====================================================================
 
         # TSCM Baselines - Environment snapshots for comparison
-        conn.execute('''
-            CREATE TABLE IF NOT EXISTS tscm_baselines (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                name TEXT NOT NULL,
-                location TEXT,
-                description TEXT,
-                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                wifi_networks TEXT,
-                bt_devices TEXT,
-                rf_frequencies TEXT,
-                gps_coords TEXT,
-                is_active BOOLEAN DEFAULT 0
-            )
-        ''')
+        conn.execute('''
+            CREATE TABLE IF NOT EXISTS tscm_baselines (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT NOT NULL,
+                location TEXT,
+                description TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                wifi_networks TEXT,
+                wifi_clients TEXT,
+                bt_devices TEXT,
+                rf_frequencies TEXT,
+                gps_coords TEXT,
+                is_active BOOLEAN DEFAULT 0
+            )
+        ''')
+
+        # Ensure new columns exist for older databases
+        try:
+            columns = {row['name'] for row in conn.execute("PRAGMA table_info(tscm_baselines)")}
+            if 'wifi_clients' not in columns:
+                conn.execute('ALTER TABLE tscm_baselines ADD COLUMN wifi_clients TEXT')
+        except Exception as e:
+            logger.debug(f"Schema update skipped for tscm_baselines: {e}")
 
         # TSCM Sweeps - Individual sweep sessions
         conn.execute('''
@@ -685,15 +740,16 @@ def get_correlations(min_confidence: float = 0.5) -> list[dict]:
 # TSCM Functions
 # =============================================================================
 
-def create_tscm_baseline(
-    name: str,
-    location: str | None = None,
-    description: str | None = None,
-    wifi_networks: list | None = None,
-    bt_devices: list | None = None,
-    rf_frequencies: list | None = None,
-    gps_coords: dict | None = None
-) -> int:
+def create_tscm_baseline(
+    name: str,
+    location: str | None = None,
+    description: str | None = None,
+    wifi_networks: list | None = None,
+    wifi_clients: list | None = None,
+    bt_devices: list | None = None,
+    rf_frequencies: list | None = None,
+    gps_coords: dict | None = None
+) -> int:
     """
     Create a new TSCM baseline.
 
@@ -701,19 +757,20 @@ def create_tscm_baseline(
         The ID of the created baseline
     """
     with get_db() as conn:
-        cursor = conn.execute('''
-            INSERT INTO tscm_baselines
-            (name, location, description, wifi_networks, bt_devices, rf_frequencies, gps_coords)
-            VALUES (?, ?, ?, ?, ?, ?, ?)
-        ''', (
-            name,
-            location,
-            description,
-            json.dumps(wifi_networks) if wifi_networks else None,
-            json.dumps(bt_devices) if bt_devices else None,
-            json.dumps(rf_frequencies) if rf_frequencies else None,
-            json.dumps(gps_coords) if gps_coords else None
-        ))
+        cursor = conn.execute('''
+            INSERT INTO tscm_baselines
+            (name, location, description, wifi_networks, wifi_clients, bt_devices, rf_frequencies, gps_coords)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+        ''', (
+            name,
+            location,
+            description,
+            json.dumps(wifi_networks) if wifi_networks else None,
+            json.dumps(wifi_clients) if wifi_clients else None,
+            json.dumps(bt_devices) if bt_devices else None,
+            json.dumps(rf_frequencies) if rf_frequencies else None,
+            json.dumps(gps_coords) if gps_coords else None
+        ))
         return cursor.lastrowid
 
 
@@ -728,18 +785,19 @@ def get_tscm_baseline(baseline_id: int) -> dict | None:
         if row is None:
             return None
 
-        return {
-            'id': row['id'],
-            'name': row['name'],
-            'location': row['location'],
-            'description': row['description'],
-            'created_at': row['created_at'],
-            'wifi_networks': json.loads(row['wifi_networks']) if row['wifi_networks'] else [],
-            'bt_devices': json.loads(row['bt_devices']) if row['bt_devices'] else [],
-            'rf_frequencies': json.loads(row['rf_frequencies']) if row['rf_frequencies'] else [],
-            'gps_coords': json.loads(row['gps_coords']) if row['gps_coords'] else None,
-            'is_active': bool(row['is_active'])
-        }
+        return {
+            'id': row['id'],
+            'name': row['name'],
+            'location': row['location'],
+            'description': row['description'],
+            'created_at': row['created_at'],
+            'wifi_networks': json.loads(row['wifi_networks']) if row['wifi_networks'] else [],
+            'wifi_clients': json.loads(row['wifi_clients']) if row['wifi_clients'] else [],
+            'bt_devices': json.loads(row['bt_devices']) if row['bt_devices'] else [],
+            'rf_frequencies': json.loads(row['rf_frequencies']) if row['rf_frequencies'] else [],
+            'gps_coords': json.loads(row['gps_coords']) if row['gps_coords'] else None,
+            'is_active': bool(row['is_active'])
+        }
 
 
 def get_all_tscm_baselines() -> list[dict]:
@@ -781,19 +839,23 @@ def set_active_tscm_baseline(baseline_id: int) -> bool:
         return cursor.rowcount > 0
 
 
-def update_tscm_baseline(
-    baseline_id: int,
-    wifi_networks: list | None = None,
-    bt_devices: list | None = None,
-    rf_frequencies: list | None = None
-) -> bool:
+def update_tscm_baseline(
+    baseline_id: int,
+    wifi_networks: list | None = None,
+    wifi_clients: list | None = None,
+    bt_devices: list | None = None,
+    rf_frequencies: list | None = None
+) -> bool:
     """Update baseline device lists."""
     updates = []
     params = []
 
-    if wifi_networks is not None:
-        updates.append('wifi_networks = ?')
-        params.append(json.dumps(wifi_networks))
+    if wifi_networks is not None:
+        updates.append('wifi_networks = ?')
+        params.append(json.dumps(wifi_networks))
+    if wifi_clients is not None:
+        updates.append('wifi_clients = ?')
+        params.append(json.dumps(wifi_clients))
     if bt_devices is not None:
         updates.append('bt_devices = ?')
         params.append(json.dumps(bt_devices))