from sh import iw
from scanf import scanf

class iw_scan:
    __interface = None
    __raw_unparsed = []
    __raw = None

    def __init__(self, interface='wlan0'):
        self.__interface = interface

    def scan(self):
        self.__raw_unparsed = []
        self.__raw = iw('{}'.format(self.__interface), 'scan')
        if self.__raw.exit_code == 0:
            self.__raw_unparsed = self.__raw.stdout.decode("utf-8").split('\n')
            return True
        return False

    def getRawData (self):
        return self.__raw_unparsed

    def getInterface(self):
        return self.__interface

    def getRawObject(self):
        return self.__raw


class iwScan:
    __iw_scan = None
    __station_list = {}

    def __init__(self, interface='wlan0'):
        self.__iw_scan = iw_scan(interface)

    def scan(self):
        if self.__iw_scan.scan():
            self.__parse_scan(self.__iw_scan.getRawData())
            return True
        return False

    def getLastError(self):
        rObject = self.__iw_scan.getRawObject()
        if rObject.exit_code != 0:
            return rObject.stderr
        return ''

    def __Parse_BBS(self, line):
        data = {}
        res = scanf("BSS %s(on %s) -- %s", line)
        if res == None:
            res = scanf("BSS %s(on %s)", line)
        data["MAC"] = res[0]
        data["DEV"] = res[1]
        if len(res) == 3:
            data["ASSOCIATED"] = 'YES'
        else:
            data["ASSOCIATED"] = 'NO'
        return data

    def __Parse_T(self, id, OnNotFoundAttr ,line):
        data = {}
        res = scanf("{}: %s".format(id), line)
        if res == None:
            data[id.upper()] = OnNotFoundAttr
        else:
            data[id.upper()] = res[0]
        return data

    def __Parse_X(self, line):
        data = {}
        res = scanf("DS Parameter set: channel %s", line)
        if res == None:
            data['CHANNEL'] = '-1'
        else:
            data['CHANNEL'] = res[0]
        return data

    def __block_stationlist(self, rawdata):
        list = {}
        count = 0
        for line in rawdata:
            idx = line.find('BSS')
            if idx != -1 and idx == 0:
                if len(list) > 0:
                    count += 1
                    self.__station_list[count] = list
                list = self.__Parse_BBS(line)
            elif line.find('SSID:') != -1:
                list = {**list, **self.__Parse_T('SSID', 'HIDDEN', line)}
            elif line.find('freq:') != -1:
                list = {**list, **self.__Parse_T('freq', '', line)}
            elif line.find('signal:') != -1:
                list = {**list, **self.__Parse_T('signal', '', line)}
            elif line.find('DS Parameter set:') != -1:
                list = {**list, **self.__Parse_X(line)}
        if count > 0:
            count += 1
            self.__station_list[count] = list

    def __parse_scan(self, rawData):
        self.__station_list = {}
        self.__block_stationlist(rawData)
        #print('{}'.format(self.__station_list))
        #for i in self.__station_list:
        #   print(self.__station_list[i])

    def findConnected (self):
        for i in self.__station_list:
            st = self.__station_list[i]
            if st.get('ASSOCIATED', 'NO') == 'YES':
                return True, self.__station_list[i]
        return False, None

    def findBySSID(self, ssid):
        for i in self.__station_list:
            st = self.__station_list[i]
            if st.get('SSID', '') == ssid:
                return True, self.__station_list[i]
        return False, None

    def getStations(self):
        return self.__station_list