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

utils/tscm/baseline.py

@@ -26,12 +26,13 @@ class BaselineRecorder:
     Records and manages TSCM environment baselines.
     """
 
-    def __init__(self):
-        self.recording = False
-        self.current_baseline_id: int | None = None
-        self.wifi_networks: dict[str, dict] = {}  # BSSID -> network info
-        self.bt_devices: dict[str, dict] = {}  # MAC -> device info
-        self.rf_frequencies: dict[float, dict] = {}  # Frequency -> signal info
+    def __init__(self):
+        self.recording = False
+        self.current_baseline_id: int | None = None
+        self.wifi_networks: dict[str, dict] = {}  # BSSID -> network info
+        self.wifi_clients: dict[str, dict] = {}  # MAC -> client info
+        self.bt_devices: dict[str, dict] = {}  # MAC -> device info
+        self.rf_frequencies: dict[float, dict] = {}  # Frequency -> signal info
 
     def start_recording(
         self,
@@ -50,10 +51,11 @@ class BaselineRecorder:
         Returns:
             Baseline ID
         """
-        self.recording = True
-        self.wifi_networks = {}
-        self.bt_devices = {}
-        self.rf_frequencies = {}
+        self.recording = True
+        self.wifi_networks = {}
+        self.wifi_clients = {}
+        self.bt_devices = {}
+        self.rf_frequencies = {}
 
         # Create baseline in database
         self.current_baseline_id = create_tscm_baseline(
@@ -78,24 +80,27 @@ class BaselineRecorder:
         self.recording = False
 
         # Convert to lists for storage
-        wifi_list = list(self.wifi_networks.values())
-        bt_list = list(self.bt_devices.values())
-        rf_list = list(self.rf_frequencies.values())
+        wifi_list = list(self.wifi_networks.values())
+        wifi_client_list = list(self.wifi_clients.values())
+        bt_list = list(self.bt_devices.values())
+        rf_list = list(self.rf_frequencies.values())
 
         # Update database
-        update_tscm_baseline(
-            self.current_baseline_id,
-            wifi_networks=wifi_list,
-            bt_devices=bt_list,
-            rf_frequencies=rf_list
-        )
+        update_tscm_baseline(
+            self.current_baseline_id,
+            wifi_networks=wifi_list,
+            wifi_clients=wifi_client_list,
+            bt_devices=bt_list,
+            rf_frequencies=rf_list
+        )
 
-        summary = {
-            'baseline_id': self.current_baseline_id,
-            'wifi_count': len(wifi_list),
-            'bt_count': len(bt_list),
-            'rf_count': len(rf_list),
-        }
+        summary = {
+            'baseline_id': self.current_baseline_id,
+            'wifi_count': len(wifi_list),
+            'wifi_client_count': len(wifi_client_list),
+            'bt_count': len(bt_list),
+            'rf_count': len(rf_list),
+        }
 
         logger.info(
             f"Baseline recording complete: {summary['wifi_count']} WiFi, "
@@ -135,8 +140,8 @@ class BaselineRecorder:
                 'last_seen': datetime.now().isoformat(),
             }
 
-    def add_bt_device(self, device: dict) -> None:
-        """Add a Bluetooth device to the current baseline."""
+    def add_bt_device(self, device: dict) -> None:
+        """Add a Bluetooth device to the current baseline."""
         if not self.recording:
             return
 
@@ -150,7 +155,7 @@ class BaselineRecorder:
                 'rssi': device.get('rssi', self.bt_devices[mac].get('rssi')),
             })
         else:
-            self.bt_devices[mac] = {
+            self.bt_devices[mac] = {
                 'mac': mac,
                 'name': device.get('name', ''),
                 'rssi': device.get('rssi', device.get('signal')),
@@ -158,10 +163,37 @@ class BaselineRecorder:
                 'type': device.get('type', ''),
                 'first_seen': datetime.now().isoformat(),
                 'last_seen': datetime.now().isoformat(),
-            }
-
-    def add_rf_signal(self, signal: dict) -> None:
-        """Add an RF signal to the current baseline."""
+            }
+
+    def add_wifi_client(self, client: dict) -> None:
+        """Add a WiFi client to the current baseline."""
+        if not self.recording:
+            return
+
+        mac = client.get('mac', client.get('address', '')).upper()
+        if not mac:
+            return
+
+        if mac in self.wifi_clients:
+            self.wifi_clients[mac].update({
+                'last_seen': datetime.now().isoformat(),
+                'rssi': client.get('rssi', self.wifi_clients[mac].get('rssi')),
+                'associated_bssid': client.get('associated_bssid', self.wifi_clients[mac].get('associated_bssid')),
+            })
+        else:
+            self.wifi_clients[mac] = {
+                'mac': mac,
+                'vendor': client.get('vendor', ''),
+                'rssi': client.get('rssi'),
+                'associated_bssid': client.get('associated_bssid'),
+                'probed_ssids': client.get('probed_ssids', []),
+                'probe_count': client.get('probe_count', len(client.get('probed_ssids', []))),
+                'first_seen': datetime.now().isoformat(),
+                'last_seen': datetime.now().isoformat(),
+            }
+
+    def add_rf_signal(self, signal: dict) -> None:
+        """Add an RF signal to the current baseline."""
         if not self.recording:
             return
 
@@ -191,15 +223,16 @@ class BaselineRecorder:
                 'hit_count': 1,
             }
 
-    def get_recording_status(self) -> dict:
-        """Get current recording status and counts."""
-        return {
-            'recording': self.recording,
-            'baseline_id': self.current_baseline_id,
-            'wifi_count': len(self.wifi_networks),
-            'bt_count': len(self.bt_devices),
-            'rf_count': len(self.rf_frequencies),
-        }
+    def get_recording_status(self) -> dict:
+        """Get current recording status and counts."""
+        return {
+            'recording': self.recording,
+            'baseline_id': self.current_baseline_id,
+            'wifi_count': len(self.wifi_networks),
+            'wifi_client_count': len(self.wifi_clients),
+            'bt_count': len(self.bt_devices),
+            'rf_count': len(self.rf_frequencies),
+        }
 
 
 class BaselineComparator:
@@ -220,11 +253,16 @@ class BaselineComparator:
             for d in baseline.get('wifi_networks', [])
             if d.get('bssid') or d.get('mac')
         }
-        self.baseline_bt = {
-            d.get('mac', d.get('address', '')).upper(): d
-            for d in baseline.get('bt_devices', [])
-            if d.get('mac') or d.get('address')
-        }
+        self.baseline_bt = {
+            d.get('mac', d.get('address', '')).upper(): d
+            for d in baseline.get('bt_devices', [])
+            if d.get('mac') or d.get('address')
+        }
+        self.baseline_wifi_clients = {
+            d.get('mac', d.get('address', '')).upper(): d
+            for d in baseline.get('wifi_clients', [])
+            if d.get('mac') or d.get('address')
+        }
         self.baseline_rf = {
             round(d.get('frequency', 0), 1): d
             for d in baseline.get('rf_frequencies', [])
@@ -269,8 +307,8 @@ class BaselineComparator:
             'matching_count': len(matching_devices),
         }
 
-    def compare_bluetooth(self, current_devices: list[dict]) -> dict:
-        """Compare current Bluetooth devices against baseline."""
+    def compare_bluetooth(self, current_devices: list[dict]) -> dict:
+        """Compare current Bluetooth devices against baseline."""
         current_macs = {
             d.get('mac', d.get('address', '')).upper(): d
             for d in current_devices
@@ -291,14 +329,45 @@ class BaselineComparator:
             if mac not in current_macs:
                 missing_devices.append(device)
 
-        return {
-            'new': new_devices,
-            'missing': missing_devices,
-            'matching': matching_devices,
-            'new_count': len(new_devices),
-            'missing_count': len(missing_devices),
-            'matching_count': len(matching_devices),
-        }
+        return {
+            'new': new_devices,
+            'missing': missing_devices,
+            'matching': matching_devices,
+            'new_count': len(new_devices),
+            'missing_count': len(missing_devices),
+            'matching_count': len(matching_devices),
+        }
+
+    def compare_wifi_clients(self, current_devices: list[dict]) -> dict:
+        """Compare current WiFi clients against baseline."""
+        current_macs = {
+            d.get('mac', d.get('address', '')).upper(): d
+            for d in current_devices
+            if d.get('mac') or d.get('address')
+        }
+
+        new_devices = []
+        missing_devices = []
+        matching_devices = []
+
+        for mac, device in current_macs.items():
+            if mac not in self.baseline_wifi_clients:
+                new_devices.append(device)
+            else:
+                matching_devices.append(device)
+
+        for mac, device in self.baseline_wifi_clients.items():
+            if mac not in current_macs:
+                missing_devices.append(device)
+
+        return {
+            'new': new_devices,
+            'missing': missing_devices,
+            'matching': matching_devices,
+            'new_count': len(new_devices),
+            'missing_count': len(missing_devices),
+            'matching_count': len(matching_devices),
+        }
 
     def compare_rf(self, current_signals: list[dict]) -> dict:
         """Compare current RF signals against baseline."""
@@ -331,35 +400,42 @@ class BaselineComparator:
             'matching_count': len(matching_signals),
         }
 
-    def compare_all(
-        self,
-        wifi_devices: list[dict] | None = None,
-        bt_devices: list[dict] | None = None,
-        rf_signals: list[dict] | None = None
-    ) -> dict:
+    def compare_all(
+        self,
+        wifi_devices: list[dict] | None = None,
+        wifi_clients: list[dict] | None = None,
+        bt_devices: list[dict] | None = None,
+        rf_signals: list[dict] | None = None
+    ) -> dict:
         """
         Compare all current data against baseline.
 
         Returns:
             Dict with comparison results for each category
         """
-        results = {
-            'wifi': None,
-            'bluetooth': None,
-            'rf': None,
-            'total_new': 0,
-            'total_missing': 0,
-        }
+        results = {
+            'wifi': None,
+            'wifi_clients': None,
+            'bluetooth': None,
+            'rf': None,
+            'total_new': 0,
+            'total_missing': 0,
+        }
 
-        if wifi_devices is not None:
-            results['wifi'] = self.compare_wifi(wifi_devices)
-            results['total_new'] += results['wifi']['new_count']
-            results['total_missing'] += results['wifi']['missing_count']
-
-        if bt_devices is not None:
-            results['bluetooth'] = self.compare_bluetooth(bt_devices)
-            results['total_new'] += results['bluetooth']['new_count']
-            results['total_missing'] += results['bluetooth']['missing_count']
+        if wifi_devices is not None:
+            results['wifi'] = self.compare_wifi(wifi_devices)
+            results['total_new'] += results['wifi']['new_count']
+            results['total_missing'] += results['wifi']['missing_count']
+
+        if wifi_clients is not None:
+            results['wifi_clients'] = self.compare_wifi_clients(wifi_clients)
+            results['total_new'] += results['wifi_clients']['new_count']
+            results['total_missing'] += results['wifi_clients']['missing_count']
+
+        if bt_devices is not None:
+            results['bluetooth'] = self.compare_bluetooth(bt_devices)
+            results['total_new'] += results['bluetooth']['new_count']
+            results['total_missing'] += results['bluetooth']['missing_count']
 
         if rf_signals is not None:
             results['rf'] = self.compare_rf(rf_signals)
@@ -369,11 +445,12 @@ class BaselineComparator:
         return results
 
 
-def get_comparison_for_active_baseline(
-    wifi_devices: list[dict] | None = None,
-    bt_devices: list[dict] | None = None,
-    rf_signals: list[dict] | None = None
-) -> dict | None:
+def get_comparison_for_active_baseline(
+    wifi_devices: list[dict] | None = None,
+    wifi_clients: list[dict] | None = None,
+    bt_devices: list[dict] | None = None,
+    rf_signals: list[dict] | None = None
+) -> dict | None:
     """
     Convenience function to compare against the active baseline.
 
@@ -385,4 +462,4 @@ def get_comparison_for_active_baseline(
         return None
 
     comparator = BaselineComparator(baseline)
-    return comparator.compare_all(wifi_devices, bt_devices, rf_signals)
+    return comparator.compare_all(wifi_devices, wifi_clients, bt_devices, rf_signals)