diff --git a/custom_components/mikrotik_router/helper.py b/custom_components/mikrotik_router/helper.py index 9c24298..55c517f 100644 --- a/custom_components/mikrotik_router/helper.py +++ b/custom_components/mikrotik_router/helper.py @@ -37,22 +37,24 @@ def from_entry_bool(entry, param, default=False, reverse=False): # --------------------------- # from_list # --------------------------- -async def from_list(data=None, source=None, key=None, vals=[], ensure_vals=[]): +async def from_list(data=None, source=None, key=None, key_search=None, vals=[], ensure_vals=[]): if not source: return data + keymap = generate_keymap(data, key_search) + for entry in source: - if key not in entry: + uid = await get_uid(entry, key) + if keymap and key_search in entry and entry[key_search] in keymap: + uid = keymap[entry[key_search]] + + if not uid: continue - if not entry[key]: - continue - - _LOGGER.debug("Processing entry {}, entry {}".format(source, entry)) - uid = entry[key] if uid not in data: data[uid] = {} + _LOGGER.debug("Processing entry {}, entry {}".format(source, entry)) for val in vals: _name = val['name'] _type = val['type'] if 'type' in val else 'str' @@ -75,3 +77,32 @@ async def from_list(data=None, source=None, key=None, vals=[], ensure_vals=[]): data[uid][val['name']] = _default return data + + +# --------------------------- +# get_uid +# --------------------------- +async def get_uid(entry, key): + if key not in entry: + return False + + if not entry[key]: + return False + + return entry[key] + + +# --------------------------- +# generate_keymap +# --------------------------- +async def generate_keymap(data, key_search): + if not key_search: + return None + + for uid in data: + if key_search not in uid: + continue + + keymap[data[uid]['name']] = data[uid]['default-name'] + + return keymap diff --git a/custom_components/mikrotik_router/mikrotik_controller.py b/custom_components/mikrotik_router/mikrotik_controller.py index 0f84977..18ad54e 100644 --- a/custom_components/mikrotik_router/mikrotik_controller.py +++ b/custom_components/mikrotik_router/mikrotik_controller.py @@ -33,7 +33,6 @@ class MikrotikControllerData(): self.data = {'routerboard': {}, 'resource': {}, 'interface': {}, - 'interface_map': {}, 'arp': {}, 'nat': {}, 'fw-update': {}, @@ -191,15 +190,6 @@ class MikrotikControllerData(): ] ) - interface_list = "" - for uid in self.data['interface']: - self.data['interface_map'][self.data['interface'][uid]['name']] = self.data['interface'][uid]['default-name'] - - if interface_list: - interface_list += "," - - interface_list += self.data['interface'][uid]['name'] - await self.get_interface_traffic(interface_list) return @@ -207,17 +197,23 @@ class MikrotikControllerData(): # get_interface_traffic # --------------------------- async def get_interface_traffic(self, interface_list): - data = await self.hass.async_add_executor_job(self.api.get_traffic, interface_list) - for entry in data: - iface_name = from_entry(entry, 'name') - if iface_name not in self.data['interface_map']: - continue + """Get traffic for all interfaces from Mikrotik""" + interface_list = "" + for uid in self.data['interface']: + if interface_list: + interface_list += "," - _LOGGER.debug("Processing entry {}, entry {}".format("/interface/monitor-traffic", entry)) - uid = self.data['interface_map'][iface_name] - self.data['interface'][uid]['rx-bits-per-second'] = from_entry(entry, 'rx-bits-per-second', default=0) - self.data['interface'][uid]['tx-bits-per-second'] = from_entry(entry, 'tx-bits-per-second', default=0) + interface_list += self.data['interface'][uid]['name'] + self.data['interface'] = await from_list( + data=self.data['interface'], + source=await self.hass.async_add_executor_job(self.api.get_traffic, interface_list), + key_search='name', + vals=[ + {'name': 'rx-bits-per-second', 'default': 0}, + {'name': 'tx-bits-per-second', 'default': 0}, + ] + ) return # ---------------------------