import unittest

import wifi

LINK = """Connected to 9c:3d:cf:12:34:56 (on wlan0)
\tSSID: Home:Net
\tfreq: 5180
\tRX: 123456 bytes (789 packets)
\tTX: 6543 bytes (98 packets)
\tsignal: -52 dBm
\trx bitrate: 6.0 MBit/s
\ttx bitrate: 866.7 MBit/s VHT-MCS 9 80MHz short GI VHT-NSS 2

\tbss flags:\tshort-slot-time
\tdtim period:\t1
\tbeacon int:\t100
"""

STATION = """Station 9c:3d:cf:12:34:56 (on wlan0)
\tinactive time:\t20 ms
\trx bytes:\t123456
\trx packets:\t789
\ttx bytes:\t6543
\ttx packets:\t98
\ttx retries:\t42
\ttx failed:\t3
\tbeacon loss:\t1
\tbeacon rx:\t1000
\trx drop misc:\t5
\tsignal:  \t-52 [-54, -58] dBm
\tsignal avg:\t-53 [-55, -57] dBm
\ttx bitrate:\t866.7 MBit/s VHT-MCS 9 80MHz short GI VHT-NSS 2
\trx bitrate:\t6.0 MBit/s
\tauthorized:\tyes
\tconnected time:\t12345 seconds
"""

SCAN = """BSS 9c:3d:cf:12:34:56(on wlan0) -- associated
\tlast seen: 1234.567s [boottime]
\tTSF: 1 usec (0d, 00:00:00)
\tfreq: 5180.0
\tbeacon interval: 100 TUs
\tcapability: ESS Privacy ShortSlotTime (0x0411)
\tsignal: -52.00 dBm
\tlast seen: 0 ms ago
\tInformation elements from Probe Response frame:
\tSSID: Home:Net
\tSupported rates: 6.0* 9.0 12.0* 18.0 24.0* 36.0 48.0 54.0
\tRSN:\t * Version: 1
\t\t * Group cipher: CCMP
\t\t * Pairwise ciphers: CCMP
\t\t * Authentication suites: PSK SAE
\t\t * Capabilities: 1-PTKSA-RC 1-GTKSA-RC MFP-capable (0x0080)
BSS 00:11:22:33:44:55(on wlan0)
\tfreq: 2437
\tcapability: ESS (0x0401)
\tsignal: -71.00 dBm
\tSSID: Cafe
\tDS Parameter set: channel 6
BSS 66:77:88:99:aa:bb(on wlan0)
\tfreq: 2412
\tcapability: ESS Privacy (0x0411)
\tsignal: -80.00 dBm
\tSSID:
\tWPA:\t * Version: 1
\t\t * Group cipher: TKIP
\t\t * Pairwise ciphers: TKIP
\t\t * Authentication suites: PSK
BSS cc:dd:ee:ff:00:11(on wlan0)
\tfreq: 2462
\tcapability: ESS Privacy (0x0411)
\tsignal: -85.00 dBm
\tSSID: Legacy
"""


class ParseLinkTest(unittest.TestCase):
    def test_connected(self):
        info = wifi.parse_link(LINK)
        self.assertTrue(info["connected"])
        self.assertEqual(info["bssid"], "9c:3d:cf:12:34:56")
        self.assertEqual(info["ssid"], "Home:Net")
        self.assertEqual(info["freq_mhz"], 5180)
        self.assertEqual(info["channel"], 36)
        self.assertEqual(info["band"], "5 GHz")
        self.assertEqual(info["signal_dbm"], -52)
        self.assertEqual(info["tx_bitrate_mbps"], 866.7)
        self.assertEqual(info["rx_bitrate_mbps"], 6.0)

    def test_not_connected(self):
        self.assertEqual(wifi.parse_link("Not connected.\n"), {"connected": False})
        self.assertEqual(wifi.parse_link(""), {"connected": False})


class ParseStationTest(unittest.TestCase):
    def test_counters(self):
        info = wifi.parse_station_dump(STATION)
        self.assertEqual(info["tx_retries"], 42)
        self.assertEqual(info["tx_failed"], 3)
        self.assertEqual(info["beacon_loss"], 1)
        self.assertEqual(info["connected_s"], 12345)
        self.assertEqual(info["signal_avg_dbm"], -53)

    def test_empty(self):
        self.assertEqual(wifi.parse_station_dump(""), {})


class ParseScanTest(unittest.TestCase):
    def test_networks_sorted_by_signal(self):
        aps = wifi.parse_scan(SCAN)
        self.assertEqual([ap["bssid"] for ap in aps], [
            "9c:3d:cf:12:34:56", "00:11:22:33:44:55", "66:77:88:99:aa:bb", "cc:dd:ee:ff:00:11",
        ])

    def test_fields(self):
        home, cafe, hidden, legacy = wifi.parse_scan(SCAN)
        self.assertTrue(home["associated"])
        self.assertEqual(home["ssid"], "Home:Net")
        self.assertEqual(home["freq_mhz"], 5180)
        self.assertEqual(home["channel"], 36)
        self.assertEqual(home["signal_dbm"], -52.0)
        self.assertEqual(home["security"], "WPA3/WPA2")

        self.assertFalse(cafe["associated"])
        self.assertEqual(cafe["channel"], 6)
        self.assertEqual(cafe["band"], "2.4 GHz")
        self.assertEqual(cafe["security"], "Open")

        self.assertEqual(hidden["ssid"], "")
        self.assertEqual(hidden["channel"], 1)
        self.assertEqual(hidden["security"], "WPA")

        self.assertEqual(legacy["channel"], 11)
        self.assertEqual(legacy["security"], "WEP")

    def test_no_internal_keys_leak(self):
        for ap in wifi.parse_scan(SCAN):
            self.assertFalse([key for key in ap if key.startswith("_")])

    def test_empty(self):
        self.assertEqual(wifi.parse_scan(""), [])


class ChannelTest(unittest.TestCase):
    def test_mapping(self):
        self.assertEqual(wifi.freq_to_channel(2412), 1)
        self.assertEqual(wifi.freq_to_channel(2484), 14)
        self.assertEqual(wifi.freq_to_channel(5745), 149)
        self.assertEqual(wifi.freq_to_channel(5955), 1)
        self.assertIsNone(wifi.freq_to_channel(None))
        self.assertIsNone(wifi.freq_to_channel(1000))


if __name__ == "__main__":
    unittest.main()