"""
Client Capability Analysis Module
Detects client WiFi generations, vendor identification, capability mismatches, and PHY rates
"""

from scapy.all import *
from scapy.layers.dot11 import *
from collections import defaultdict, Counter


def detect_client_capabilities(packets):
    """Analyze client capabilities, PHY rates, and vendor info"""
    capabilities = {
        'clients': {},
        'capability_mismatches': [],
        'phy_rate_distribution': {},
        'vendor_analysis': {},
        'wifi_generation_by_client': {},
        'not_detected_reason': {}
    }

    # OUI database (top vendors) - first 3 bytes of MAC
    oui_db = {
        '00:50:f2': 'Microsoft',
        '00:03:93': 'Apple',
        'd8:9e:f3': 'Apple',
        'f8:ff:c2': 'Apple',
        '00:23:12': 'Apple',
        '00:1f:f3': 'Apple',
        '00:26:bb': 'Apple',
        'ac:de:48': 'Apple',
        '00:0f:b5': 'Netgear',
        '00:09:5b': 'Netgear',
        '00:1e:2a': 'Cisco',
        '00:40:96': 'Cisco',
        '00:0d:3a': 'Intel',
        '00:12:f0': 'Intel',
        '00:1b:21': 'Intel',
        'b4:b6:76': 'Intel',
        '00:21:6a': 'Intel',
        '00:24:d7': 'Intel',
        '48:89:e7': 'Intel',
        '00:0c:29': 'VMware',
        '00:50:56': 'VMware',
        '00:05:69': 'VMware',
        '00:1c:14': 'VMware',
        '2c:ea:7f': 'Samsung',
        '00:13:77': 'Samsung',
        '00:16:32': 'Samsung',
        '00:1d:25': 'Samsung',
        '00:14:22': 'Dell',
        '00:1a:a0': 'Dell',
        '00:21:70': 'Dell',
        'b8:ca:3a': 'Dell',
        '18:03:73': 'Dell',
        '3c:46:d8': 'TP-Link',
        'f4:f2:6d': 'TP-Link',
        '00:0c:43': 'Ralink',
        '00:90:4c': 'Epigram (Broadcom)',
        '00:10:18': 'Broadcom',
        '00:05:4e': 'Brocade',
    }

    client_info = defaultdict(lambda: {
        'wifi_gen': set(),
        'phy_rates': [],
        'capabilities': set(),
        'packet_count': 0,
        'vendor': 'Unknown'
    })

    ap_capabilities = {}  # Track AP WiFi generation capabilities

    for idx, pkt in enumerate(packets, 1):
        try:
            # Extract PHY rate from RadioTap
            if pkt.haslayer(RadioTap):
                rt = pkt[RadioTap]
                if hasattr(rt, 'Rate'):
                    rate_mbps = rt.Rate
                    if rate_mbps not in capabilities['phy_rate_distribution']:
                        capabilities['phy_rate_distribution'][rate_mbps] = 0
                    capabilities['phy_rate_distribution'][rate_mbps] += 1

            if pkt.haslayer(Dot11):
                dot11 = pkt[Dot11]
                client_mac = None
                is_from_ap = False

                # Identify client MAC (addr2 for frames from client)
                if dot11.type == 0:  # Management
                    if pkt.haslayer(Dot11AssoReq) or pkt.haslayer(Dot11ProbeReq):
                        client_mac = dot11.addr2
                    elif pkt.haslayer(Dot11Beacon):
                        is_from_ap = True
                        client_mac = dot11.addr3  # BSSID for AP capability tracking
                elif dot11.type == 2:  # Data
                    # FromDS=0, ToDS=1: from client to AP
                    if (dot11.FCfield & 0x01) and not (dot11.FCfield & 0x02):
                        client_mac = dot11.addr2

            # Analyze capabilities from management frames
            if pkt.haslayer(Dot11Beacon) or pkt.haslayer(Dot11ProbeResp) or pkt.haslayer(Dot11AssoReq) or pkt.haslayer(Dot11ProbeReq):
                mac = None
                if pkt.haslayer(Dot11):
                    if pkt.haslayer(Dot11Beacon):
                        mac = pkt[Dot11].addr3  # BSSID
                        is_from_ap = True
                    elif pkt.haslayer(Dot11AssoReq) or pkt.haslayer(Dot11ProbeReq):
                        mac = pkt[Dot11].addr2  # Client
                        is_from_ap = False
                    elif pkt.haslayer(Dot11ProbeResp):
                        mac = pkt[Dot11].addr3  # BSSID
                        is_from_ap = True

                if mac and pkt.haslayer(Dot11Elt):
                    elt = pkt[Dot11Elt]
                    wifi_gen = set()
                    caps = set()

                    while elt:
                        # HT Capabilities (WiFi 4)
                        if elt.ID == 45:
                            wifi_gen.add('WiFi 4 (802.11n)')
                            caps.add('HT')

                        # VHT Capabilities (WiFi 5)
                        if elt.ID == 191:
                            wifi_gen.add('WiFi 5 (802.11ac)')
                            caps.add('VHT')

                        # HE Capabilities (WiFi 6)
                        if elt.ID == 255 and len(elt.info) >= 2 and elt.info[0] == 35:
                            wifi_gen.add('WiFi 6 (802.11ax)')
                            caps.add('HE')

                        # EHT Capabilities (WiFi 7)
                        if elt.ID == 255 and len(elt.info) >= 2 and elt.info[0] == 106:
                            wifi_gen.add('WiFi 7 (802.11be)')
                            caps.add('EHT')

                        if hasattr(elt, 'payload'):
                            elt = elt.payload.getlayer(Dot11Elt)
                        else:
                            break

                    if wifi_gen:
                        if is_from_ap:
                            ap_capabilities[mac] = wifi_gen
                        else:
                            client_info[mac]['wifi_gen'].update(wifi_gen)
                            client_info[mac]['capabilities'].update(caps)
                            client_info[mac]['packet_count'] += 1

                            # Vendor identification
                            oui = mac[:8].upper()  # First 3 bytes (XX:XX:XX)
                            if oui in oui_db:
                                client_info[mac]['vendor'] = oui_db[oui]

        except Exception:
            continue

    # Build client capability summary
    for mac, info in client_info.items():
        capabilities['clients'][mac] = {
            'wifi_generation': list(info['wifi_gen']),
            'capabilities': list(info['capabilities']),
            'packet_count': info['packet_count'],
            'vendor': info['vendor']
        }

        # Determine highest WiFi generation
        if 'WiFi 7 (802.11be)' in info['wifi_gen']:
            capabilities['wifi_generation_by_client'][mac] = 'WiFi 7'
        elif 'WiFi 6 (802.11ax)' in info['wifi_gen']:
            capabilities['wifi_generation_by_client'][mac] = 'WiFi 6'
        elif 'WiFi 5 (802.11ac)' in info['wifi_gen']:
            capabilities['wifi_generation_by_client'][mac] = 'WiFi 5'
        elif 'WiFi 4 (802.11n)' in info['wifi_gen']:
            capabilities['wifi_generation_by_client'][mac] = 'WiFi 4'
        else:
            capabilities['wifi_generation_by_client'][mac] = 'Legacy'

    # Detect capability mismatches (AP supports WiFi 6/7, client only WiFi 5)
    for ap_mac, ap_gen in ap_capabilities.items():
        ap_max = 'WiFi 7' if 'WiFi 7 (802.11be)' in ap_gen else \
                 'WiFi 6' if 'WiFi 6 (802.11ax)' in ap_gen else \
                 'WiFi 5' if 'WiFi 5 (802.11ac)' in ap_gen else 'WiFi 4'

        for client_mac, client_gen_str in capabilities['wifi_generation_by_client'].items():
            # Check if client is lower generation than AP
            gen_order = {'WiFi 7': 4, 'WiFi 6': 3, 'WiFi 5': 2, 'WiFi 4': 1, 'Legacy': 0}
            if gen_order.get(client_gen_str, 0) < gen_order.get(ap_max, 0):
                capabilities['capability_mismatches'].append({
                    'client': client_mac,
                    'client_gen': client_gen_str,
                    'ap': ap_mac,
                    'ap_gen': ap_max,
                    'impact': 'Client cannot utilize AP\'s advanced features'
                })

    # Vendor analysis summary
    vendor_counts = Counter()
    for mac, info in client_info.items():
        vendor_counts[info['vendor']] += 1
    capabilities['vendor_analysis'] = dict(vendor_counts)

    return capabilities
