# credits: https://gist.github.com/calebmadrigal/fdb8855a6d05c87bbb0254a1424ee582

import sys
import yaml

IDENTITIES_PATH = "identities.yaml"

def parse_wifi_map(map_path):
    with open(map_path, 'r') as f:
        wifi_map = yaml.safe_load(f)

    with open(IDENTITIES_PATH, 'r') as f:
        identities = yaml.safe_load(f)['identities']

    for identity in identities:
        print('mac = {}, name = {}'.format(identity['mac'],identity['name']))
    
    devices = set()

    filtered_devices = []

    for ssid in wifi_map:
        print('ssid = {}'.format(ssid))
        ssid_node = wifi_map[ssid]
        for bssid in ssid_node:
            print('\tbssid = {}'.format(bssid))
            bssid_node = ssid_node[bssid]
            if 'devices' in bssid_node:
                for device in bssid_node['devices']:
                    devices |= {device}
                    print('\t\tdevice = {}'.format(device))
                    for identity in identities:
                        if identity['mac'] == device:
                            filtered_devices.append(identity)

    print('\n\nSSID count: {}, Device count: {}'.format(len(wifi_map), len(devices)))
    print('\nFiltered devices:')
    print(filtered_devices)

if __name__ == '__main__':
    wifi_map_path = 'wifi_map.yaml'
    if len(sys.argv) > 1:
        wifi_map_path = sys.argv[1]
    parse_wifi_map(wifi_map_path)