diff --git a/enrich.py b/enrich.py index 8a3fda7..699abc6 100755 --- a/enrich.py +++ b/enrich.py @@ -148,6 +148,51 @@ def get_clients_on_channel(capture, ap_channel, ap_bssid): return len(clients) +def get_aps_on_channel(capture, ap_channel): + try: + ap_channel = int(ap_channel) + except ValueError: + print(f"[!] Could not parse channel number: {ap_channel}") + return 0 + + aps = set() + + for packet in capture: + try: + if not hasattr(packet, 'radiotap') or not hasattr(packet, 'wlan'): + continue + + # Check if the packet has frequency info + if not hasattr(packet.radiotap, 'channel') or not hasattr(packet.radiotap.channel, 'freq'): + continue + + packet_freq = int(packet.radiotap.channel.freq) + packet_channel = get_channel_from_freq(packet_freq) + + if packet_channel != ap_channel: + continue + + # Identify AP-advertising frames: Beacon (0x08) or Probe Response (0x05) + subtype_hex = getattr(packet.wlan, 'fc_type_subtype', None) + if subtype_hex is None: + continue + + subtype = int(subtype_hex, 16) + if subtype not in (0x08, 0x05): + continue + + bssid = getattr(packet.wlan, 'bssid', '').lower() + if bssid: + aps.add(bssid) + + except AttributeError: + continue + except Exception as e: + print(f"[!] AP scan error: {e}") + continue + + return len(aps) + def analyze_pcap(pcapng_path, start_ts, end_ts, ap_bssid, ap_channel): cap = pyshark.FileCapture( pcapng_path, @@ -173,18 +218,18 @@ def analyze_pcap(pcapng_path, start_ts, end_ts, ap_bssid, ap_channel): clients_on_ap = get_clients_on_ap(filtered_packets, ap_bssid) clients_on_channel = get_clients_on_channel(filtered_packets, ap_channel, ap_bssid) + aps_on_channel = get_aps_on_channel(filtered_packets, ap_channel) # Placeholder: Logic will be added for: - # - APsOnChannel # - CongestionScore # - AvgAPSignal # - StrongestAPSignal # - UnlinkedDevices - + finally: cap.close() - return clients_on_ap, clients_on_channel, 0, None, None, None, 0 + return clients_on_ap, clients_on_channel, aps_on_channel, None, None, None, 0 def main():