mirror.MikroTik/tomaae.homeassistant-mikrotik_router
1
0
Fork 0
mirror of https://github.com/tomaae/homeassistant-mikrotik_router.git synced 2025-07-29 15:04:31 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

reformatted using black

Browse source
This commit is contained in:
tomaae 2020-04-11 05:45:36 +02:00
parent 8f89051be6
commit 70cdb93eb3
8 changed files with 319 additions and 222 deletions
Download patch file Download diff file Expand all files Collapse all files

12
custom_components/mikrotik_router/__init__.py
View file

@ -44,13 +44,11 @@ async def async_setup_entry(hass, config_entry):
)
hass.async_create_task(
hass.config_entries.async_forward_entry_setup(config_entry,
"binary_sensor")
hass.config_entries.async_forward_entry_setup(config_entry, "binary_sensor")
)
hass.async_create_task(
hass.config_entries.async_forward_entry_setup(config_entry,
"device_tracker")
hass.config_entries.async_forward_entry_setup(config_entry, "device_tracker")
)
hass.async_create_task(
@ -76,10 +74,8 @@ async def async_unload_entry(hass, config_entry):
"""Unload a config entry."""
controller = hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id]
await hass.config_entries.async_forward_entry_unload(config_entry, "sensor")
await hass.config_entries.async_forward_entry_unload(config_entry,
"binary_sensor")
await hass.config_entries.async_forward_entry_unload(config_entry,
"device_tracker")
await hass.config_entries.async_forward_entry_unload(config_entry, "binary_sensor")
await hass.config_entries.async_forward_entry_unload(config_entry, "device_tracker")
await hass.config_entries.async_forward_entry_unload(config_entry, "switch")
await controller.async_reset()
hass.data[DOMAIN][DATA_CLIENT].pop(config_entry.entry_id)

17
custom_components/mikrotik_router/config_flow.py
View file

@ -50,8 +50,7 @@ _LOGGER = logging.getLogger(__name__)
def configured_instances(hass):
"""Return a set of configured instances."""
return set(
entry.data[CONF_NAME] for entry in
hass.config_entries.async_entries(DOMAIN)
entry.data[CONF_NAME] for entry in hass.config_entries.async_entries(DOMAIN)
)
@ -91,7 +90,7 @@ class MikrotikControllerConfigFlow(ConfigFlow, domain=DOMAIN):
username=user_input[CONF_USERNAME],
password=user_input[CONF_PASSWORD],
port=user_input[CONF_PORT],
use_ssl=user_input[CONF_SSL]
use_ssl=user_input[CONF_SSL],
)
if not api.connect():
errors[CONF_HOST] = api.error
@ -99,14 +98,10 @@ class MikrotikControllerConfigFlow(ConfigFlow, domain=DOMAIN):
# Save instance
if not errors:
return self.async_create_entry(
title=user_input[CONF_NAME],
data=user_input
title=user_input[CONF_NAME], data=user_input
)
return self._show_config_form(
user_input=user_input,
errors=errors,
)
return self._show_config_form(user_input=user_input, errors=errors)
return self._show_config_form(
user_input={
@ -117,7 +112,7 @@ class MikrotikControllerConfigFlow(ConfigFlow, domain=DOMAIN):
CONF_PORT: DEFAULT_PORT,
CONF_SSL: DEFAULT_SSL,
},
errors=errors
errors=errors,
)
# ---------------------------
@ -195,7 +190,7 @@ class MikrotikControllerOptionsFlowHandler(OptionsFlow):
default=self.config_entry.options.get(
CONF_TRACK_HOSTS_TIMEOUT, DEFAULT_TRACK_HOST_TIMEOUT
),
): int
): int,
}
),
)

33
custom_components/mikrotik_router/device_tracker.py
View file

@ -88,7 +88,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback
def update_controller():
"""Update the values of the controller."""
update_items(inst, config_entry, mikrotik_controller, async_add_entities, tracked)
update_items(
inst, config_entry, mikrotik_controller, async_add_entities, tracked
)
mikrotik_controller.listeners.append(
async_dispatcher_connect(
@ -111,10 +113,7 @@ def update_items(inst, config_entry, mikrotik_controller, async_add_entities, tr
for sid, sid_uid, sid_func in zip(
["interface", "host"],
["default-name", "mac-address"],
[
MikrotikControllerPortDeviceTracker,
MikrotikControllerHostDeviceTracker,
],
[MikrotikControllerPortDeviceTracker, MikrotikControllerHostDeviceTracker],
):
for uid in mikrotik_controller.data[sid]:
item_id = f"{inst}-{sid}-{mikrotik_controller.data[sid][uid][sid_uid]}"
@ -206,8 +205,7 @@ class MikrotikControllerPortDeviceTracker(ScannerEntity):
def device_info(self):
"""Return a port description for device registry."""
info = {
"connections": {
(CONNECTION_NETWORK_MAC, self._data["port-mac-address"])},
"connections": {(CONNECTION_NETWORK_MAC, self._data["port-mac-address"])},
"manufacturer": self._ctrl.data["resource"]["platform"],
"model": self._ctrl.data["resource"]["board-name"],
"name": f"{self._inst} {self._data['default-name']}",
@ -333,11 +331,10 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity):
def device_info(self):
"""Return a host description for device registry."""
info = {
"connections": {
(CONNECTION_NETWORK_MAC, self._data["mac-address"])},
"connections": {(CONNECTION_NETWORK_MAC, self._data["mac-address"])},
"manufacturer": self._ctrl.data["resource"]["platform"],
"model": self._ctrl.data["resource"]["board-name"],
"name": self._data['host-name'],
"name": self._data["host-name"],
}
return info
@ -349,12 +346,22 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity):
if variable in self._data:
if variable == "last-seen":
if self._data[variable]:
attributes[format_attribute(variable)] = get_age(self._data[variable])
attributes[format_attribute(variable)] = get_age(
self._data[variable]
)
else:
attributes[format_attribute(variable)] = "unknown"
else:
if self._data[variable] in ["dhcp", "dns", "capsman", "wireless", "restored"]:
attributes[format_attribute(variable)] = format_value(self._data[variable])
if self._data[variable] in [
"dhcp",
"dns",
"capsman",
"wireless",
"restored",
]:
attributes[format_attribute(variable)] = format_value(
self._data[variable]
)
else:
attributes[format_attribute(variable)] = self._data[variable]

3
custom_components/mikrotik_router/helper.py
View file

@ -266,8 +266,7 @@ def fill_vals_proc(data, uid, vals_proc) -> dict:
if _action == "combine":
if "key" in val:
tmp = _data[val["key"]]\
if val["key"] in _data else "unknown"
tmp = _data[val["key"]] if val["key"] in _data else "unknown"
if not _value:
_value = tmp
else:

342
custom_components/mikrotik_router/mikrotik_controller.py
View file

@ -32,7 +32,7 @@ from .const import (
DEFAULT_UNIT_OF_MEASUREMENT,
)
from .exceptions import ApiEntryNotFound
from .helper import from_entry, parse_api
from .helper import parse_api
from .mikrotikapi import MikrotikAPI
_LOGGER = logging.getLogger(__name__)
@ -70,7 +70,7 @@ class MikrotikControllerData:
"wireless_hosts": {},
"host": {},
"host_hass": {},
"accounting": {}
"accounting": {},
}
self.listeners = []
@ -82,7 +82,7 @@ class MikrotikControllerData:
config_entry.data[CONF_USERNAME],
config_entry.data[CONF_PASSWORD],
config_entry.data[CONF_PORT],
config_entry.data[CONF_SSL]
config_entry.data[CONF_SSL],
)
self.api_ping = MikrotikAPI(
@ -90,7 +90,7 @@ class MikrotikControllerData:
config_entry.data[CONF_USERNAME],
config_entry.data[CONF_PASSWORD],
config_entry.data[CONF_PORT],
config_entry.data[CONF_SSL]
config_entry.data[CONF_SSL],
)
self.nat_removed = {}
@ -121,7 +121,9 @@ class MikrotikControllerData:
@property
def option_track_iface_clients(self):
"""Config entry option to not track ARP."""
return self.config_entry.options.get(CONF_TRACK_IFACE_CLIENTS, DEFAULT_TRACK_IFACE_CLIENTS)
return self.config_entry.options.get(
CONF_TRACK_IFACE_CLIENTS, DEFAULT_TRACK_IFACE_CLIENTS
)
# ---------------------------
# option_track_network_hosts
@ -229,9 +231,11 @@ class MikrotikControllerData:
"""Get host data from HA entity registry"""
registry = await self.hass.helpers.entity_registry.async_get_registry()
for entity in registry.entities.values():
if entity.config_entry_id == self.config_entry.entry_id \
and entity.domain == DEVICE_TRACKER_DOMAIN \
and "-host-" in entity.unique_id:
if (
entity.config_entry_id == self.config_entry.entry_id
and entity.domain == DEVICE_TRACKER_DOMAIN
and "-host-" in entity.unique_id
):
_, mac = entity.unique_id.split("-host-", 2)
self.data["host_hass"][mac] = entity.original_name
@ -283,17 +287,30 @@ class MikrotikControllerData:
for uid, vals in self.data["host"].items():
# Add missing default values
for key, default in zip(
["address", "mac-address", "interface", "host-name", "last-seen", "available"],
[
"address",
"mac-address",
"interface",
"host-name",
"last-seen",
"available",
],
["unknown", "unknown", "unknown", "unknown", False, False],
):
if key not in self.data["host"][uid]:
self.data["host"][uid][key] = default
# Check host availability
if vals["source"] not in ["capsman", "wireless"] \
and vals["address"] != "unknown" and vals["interface"] != "unknown":
self.data["host"][uid]["available"] = \
await self.hass.async_add_executor_job(self.api_ping.arp_ping, vals["address"], vals["interface"])
if (
vals["source"] not in ["capsman", "wireless"]
and vals["address"] != "unknown"
and vals["interface"] != "unknown"
):
self.data["host"][uid][
"available"
] = await self.hass.async_add_executor_job(
self.api_ping.arp_ping, vals["address"], vals["interface"]
)
# Update last seen
if self.data["host"][uid]["available"]:
@ -409,10 +426,8 @@ class MikrotikControllerData:
uom_type, uom_div = self._get_unit_of_measurement()
for uid in self.data["interface"]:
self.data["interface"][uid][
"rx-bits-per-second-attr"] = uom_type
self.data["interface"][uid][
"tx-bits-per-second-attr"] = uom_type
self.data["interface"][uid]["rx-bits-per-second-attr"] = uom_type
self.data["interface"][uid]["tx-bits-per-second-attr"] = uom_type
self.data["interface"][uid]["rx-bits-per-second"] = round(
self.data["interface"][uid]["rx-bits-per-second"] * uom_div
)
@ -440,7 +455,9 @@ class MikrotikControllerData:
"reverse": True,
},
],
only=[{"key": "local", "value": False}],
only=[
{"key": "local", "value": False}
],
)
for uid, vals in self.data["bridge_host"].items():
@ -465,12 +482,16 @@ class MikrotikControllerData:
continue
if self.data["interface"][uid]["client-ip-address"] == "":
self.data["interface"][uid]["client-ip-address"] = arp_vals["address"]
self.data["interface"][uid]["client-ip-address"] = arp_vals[
"address"
]
else:
self.data["interface"][uid]["client-ip-address"] = "multiple"
if self.data["interface"][uid]["client-mac-address"] == "":
self.data["interface"][uid]["client-mac-address"] = arp_vals["mac-address"]
self.data["interface"][uid]["client-mac-address"] = arp_vals[
"mac-address"
]
else:
self.data["interface"][uid]["client-mac-address"] = "multiple"
@ -513,7 +534,9 @@ class MikrotikControllerData:
{"key": "dst-port"},
]
],
only=[{"key": "action", "value": "dst-nat"}],
only=[
{"key": "action", "value": "dst-nat"}
],
)
# Remove duplicate NAT entries to prevent crash
@ -530,8 +553,11 @@ class MikrotikControllerData:
for uid in nat_del:
if self.data["nat"][uid]["name"] not in self.nat_removed:
self.nat_removed[self.data["nat"][uid]["name"]] = 1
_LOGGER.error("Mikrotik %s duplicate NAT rule %s, entity will be unavailable.",
self.host, self.data["nat"][uid]["name"])
_LOGGER.error(
"Mikrotik %s duplicate NAT rule %s, entity will be unavailable.",
self.host,
self.data["nat"][uid]["name"],
)
del self.data["nat"][uid]
@ -619,8 +645,8 @@ class MikrotikControllerData:
if "status" in self.data["fw-update"]:
self.data["fw-update"]["available"] = (
True if self.data["fw-update"][
"status"] == "New version is available"
True
if self.data["fw-update"]["status"] == "New version is available"
else False
)
else:
@ -671,40 +697,52 @@ class MikrotikControllerData:
"type": "bool",
"reverse": True,
},
]
],
)
uom_type, uom_div = self._get_unit_of_measurement()
for uid, vals in self.data["queue"].items():
upload_max_limit_bps, download_max_limit_bps = \
[int(x) for x in vals["max-limit"].split('/')]
self.data["queue"][uid]["upload-max-limit"] = \
f"{round(upload_max_limit_bps * uom_div)} {uom_type}"
self.data["queue"][uid]["download-max-limit"] = \
f"{round(download_max_limit_bps * uom_div)} {uom_type}"
upload_max_limit_bps, download_max_limit_bps = [
int(x) for x in vals["max-limit"].split("/")
]
self.data["queue"][uid][
"upload-max-limit"
] = f"{round(upload_max_limit_bps * uom_div)} {uom_type}"
self.data["queue"][uid][
"download-max-limit"
] = f"{round(download_max_limit_bps * uom_div)} {uom_type}"
upload_limit_at_bps, download_limit_at_bps = \
[int(x) for x in vals["limit-at"].split('/')]
self.data["queue"][uid]["upload-limit-at"] = \
f"{round(upload_limit_at_bps * uom_div)} {uom_type}"
self.data["queue"][uid]["download-limit-at"] = \
f"{round(download_limit_at_bps * uom_div)} {uom_type}"
upload_limit_at_bps, download_limit_at_bps = [
int(x) for x in vals["limit-at"].split("/")
]
self.data["queue"][uid][
"upload-limit-at"
] = f"{round(upload_limit_at_bps * uom_div)} {uom_type}"
self.data["queue"][uid][
"download-limit-at"
] = f"{round(download_limit_at_bps * uom_div)} {uom_type}"
upload_burst_limit_bps, download_burst_limit_bps = \
[int(x) for x in vals["burst-limit"].split('/')]
self.data["queue"][uid]["upload-burst-limit"] = \
f"{round(upload_burst_limit_bps * uom_div)} {uom_type}"
self.data["queue"][uid]["download-burst-limit"] = \
f"{round(download_burst_limit_bps * uom_div)} {uom_type}"
upload_burst_limit_bps, download_burst_limit_bps = [
int(x) for x in vals["burst-limit"].split("/")
]
self.data["queue"][uid][
"upload-burst-limit"
] = f"{round(upload_burst_limit_bps * uom_div)} {uom_type}"
self.data["queue"][uid][
"download-burst-limit"
] = f"{round(download_burst_limit_bps * uom_div)} {uom_type}"
upload_burst_threshold_bps, download_burst_threshold_bps = \
[int(x) for x in vals["burst-threshold"].split('/')]
self.data["queue"][uid]["upload-burst-threshold"] = \
f"{round(upload_burst_threshold_bps * uom_div)} {uom_type}"
self.data["queue"][uid]["download-burst-threshold"] = \
f"{round(download_burst_threshold_bps * uom_div)} {uom_type}"
upload_burst_threshold_bps, download_burst_threshold_bps = [
int(x) for x in vals["burst-threshold"].split("/")
]
self.data["queue"][uid][
"upload-burst-threshold"
] = f"{round(upload_burst_threshold_bps * uom_div)} {uom_type}"
self.data["queue"][uid][
"download-burst-threshold"
] = f"{round(download_burst_threshold_bps * uom_div)} {uom_type}"
upload_burst_time, download_burst_time = vals["burst-time"].split('/')
upload_burst_time, download_burst_time = vals["burst-time"].split("/")
self.data["queue"][uid]["upload-burst-time"] = upload_burst_time
self.data["queue"][uid]["download-burst-time"] = download_burst_time
@ -724,13 +762,18 @@ class MikrotikControllerData:
],
ensure_vals=[
{"name": "bridge", "default": ""},
]
],
)
for uid, vals in self.data["arp"].items():
if vals["interface"] in self.data["bridge"] and uid in self.data["bridge_host"]:
if (
vals["interface"] in self.data["bridge"]
and uid in self.data["bridge_host"]
):
self.data["arp"][uid]["bridge"] = vals["interface"]
self.data["arp"][uid]["interface"] = self.data["bridge_host"][uid]["interface"]
self.data["arp"][uid]["interface"] = self.data["bridge_host"][uid][
"interface"
]
# ---------------------------
# get_dns
@ -768,12 +811,14 @@ class MikrotikControllerData:
ensure_vals=[
{"name": "address"},
{"name": "IPv4Network", "default": ""},
]
],
)
for uid, vals in self.data["dhcp-network"].items():
if vals["IPv4Network"] == "":
self.data["dhcp-network"][uid]["IPv4Network"] = IPv4Network(vals["address"])
self.data["dhcp-network"][uid]["IPv4Network"] = IPv4Network(
vals["address"]
)
# TODO: run only on demand
self.data["dhcp-server"] = parse_api(
@ -783,7 +828,7 @@ class MikrotikControllerData:
vals=[
{"name": "name"},
{"name": "interface", "default": ""},
]
],
)
self.data["dhcp"] = parse_api(
@ -801,12 +846,13 @@ class MikrotikControllerData:
],
ensure_vals=[
{"name": "interface"},
]
],
)
for uid in self.data["dhcp"]:
self.data["dhcp"][uid]["interface"] = \
self.data["dhcp-server"][self.data["dhcp"][uid]["server"]]["interface"]
self.data["dhcp"][uid]["interface"] = self.data["dhcp-server"][
self.data["dhcp"][uid]["server"]
]["interface"]
# ---------------------------
# get_capsman_hosts
@ -821,7 +867,7 @@ class MikrotikControllerData:
{"name": "mac-address"},
{"name": "interface", "default": "unknown"},
{"name": "ssid", "default": "unknown"},
]
],
)
# ---------------------------
@ -838,7 +884,7 @@ class MikrotikControllerData:
{"name": "interface", "default": "unknown"},
{"name": "ap", "type": "bool"},
{"name": "uptime"},
]
],
)
# ---------------------------
@ -847,8 +893,8 @@ class MikrotikControllerData:
async def async_process_host(self):
"""Get host tracking data"""
# Add hosts from CAPS-MAN
if self.support_capsman:
capsman_detected = {}
if self.support_capsman:
for uid, vals in self.data["capsman_hosts"].items():
if uid not in self.data["host"]:
self.data["host"][uid] = {}
@ -858,12 +904,15 @@ class MikrotikControllerData:
self.data["host"][uid]["available"] = True
self.data["host"][uid]["last-seen"] = utcnow()
for key in ["mac-address", "interface"]:
if key not in self.data["host"][uid] or self.data["host"][uid][key] == "unknown":
if (
key not in self.data["host"][uid]
or self.data["host"][uid][key] == "unknown"
):
self.data["host"][uid][key] = vals[key]
# Add hosts from wireless
if self.support_wireless:
wireless_detected = {}
if self.support_wireless:
for uid, vals in self.data["wireless_hosts"].items():
if vals["ap"]:
continue
@ -876,7 +925,10 @@ class MikrotikControllerData:
self.data["host"][uid]["available"] = True
self.data["host"][uid]["last-seen"] = utcnow()
for key in ["mac-address", "interface"]:
if key not in self.data["host"][uid] or self.data["host"][uid][key] == "unknown":
if (
key not in self.data["host"][uid]
or self.data["host"][uid][key] == "unknown"
):
self.data["host"][uid][key] = vals[key]
# Add hosts from DHCP
@ -885,7 +937,10 @@ class MikrotikControllerData:
self.data["host"][uid] = {}
self.data["host"][uid]["source"] = "dhcp"
for key in ["address", "mac-address", "interface"]:
if key not in self.data["host"][uid] or self.data["host"][uid][key] == "unknown":
if (
key not in self.data["host"][uid]
or self.data["host"][uid][key] == "unknown"
):
self.data["host"][uid][key] = vals[key]
# Add hosts from ARP
@ -894,7 +949,10 @@ class MikrotikControllerData:
self.data["host"][uid] = {}
self.data["host"][uid]["source"] = "arp"
for key in ["address", "mac-address", "interface"]:
if key not in self.data["host"][uid] or self.data["host"][uid][key] == "unknown":
if (
key not in self.data["host"][uid]
or self.data["host"][uid][key] == "unknown"
):
self.data["host"][uid][key] = vals[key]
# Add restored hosts from hass registry
@ -922,34 +980,59 @@ class MikrotikControllerData:
# Update IP and interface (DHCP/returned host)
if uid in self.data["dhcp"] and "." in self.data["dhcp"][uid]["address"]:
if self.data["dhcp"][uid]["address"] != self.data["host"][uid]["address"]:
self.data["host"][uid]["address"] = self.data["dhcp"][uid]["address"]
if (
self.data["dhcp"][uid]["address"]
!= self.data["host"][uid]["address"]
):
self.data["host"][uid]["address"] = self.data["dhcp"][uid][
"address"
]
if vals["source"] not in ["capsman", "wireless"]:
self.data["host"][uid]["source"] = "dhcp"
self.data["host"][uid]["interface"] = self.data["dhcp"][uid]["interface"]
self.data["host"][uid]["interface"] = self.data["dhcp"][uid][
"interface"
]
elif uid in self.data["arp"] and "." in self.data["arp"][uid]["address"] \
and self.data["arp"][uid]["address"] != self.data["host"][uid]["address"]:
elif (
uid in self.data["arp"]
and "." in self.data["arp"][uid]["address"]
and self.data["arp"][uid]["address"]
!= self.data["host"][uid]["address"]
):
self.data["host"][uid]["address"] = self.data["arp"][uid]["address"]
if vals["source"] not in ["capsman", "wireless"]:
self.data["host"][uid]["source"] = "arp"
self.data["host"][uid]["interface"] = self.data["arp"][uid]["interface"]
self.data["host"][uid]["interface"] = self.data["arp"][uid][
"interface"
]
if vals["host-name"] == "unknown":
# Resolve hostname from static DNS
if vals["address"] != "unknown":
for dns_uid, dns_vals in self.data["dns"].items():
if dns_vals["address"] == vals["address"]:
self.data["host"][uid]["host-name"] = dns_vals["name"].split('.')[0]
self.data["host"][uid]["host-name"] = dns_vals[
"name"
].split(".")[0]
break
# Resolve hostname from DHCP comment
if self.data["host"][uid]["host-name"] == "unknown" \
and uid in self.data["dhcp"] and self.data["dhcp"][uid]["comment"] != "":
self.data["host"][uid]["host-name"] = self.data["dhcp"][uid]["comment"]
if (
self.data["host"][uid]["host-name"] == "unknown"
and uid in self.data["dhcp"]
and self.data["dhcp"][uid]["comment"] != ""
):
self.data["host"][uid]["host-name"] = self.data["dhcp"][uid][
"comment"
]
# Resolve hostname from DHCP hostname
elif self.data["host"][uid]["host-name"] == "unknown" \
and uid in self.data["dhcp"] and self.data["dhcp"][uid]["host-name"] != "unknown":
self.data["host"][uid]["host-name"] = self.data["dhcp"][uid]["host-name"]
elif (
self.data["host"][uid]["host-name"] == "unknown"
and uid in self.data["dhcp"]
and self.data["dhcp"][uid]["host-name"] != "unknown"
):
self.data["host"][uid]["host-name"] = self.data["dhcp"][uid][
"host-name"
]
# Fallback to mac address for hostname
elif self.data["host"][uid]["host-name"] == "unknown":
self.data["host"][uid]["host-name"] = uid
@ -960,31 +1043,34 @@ class MikrotikControllerData:
def process_accounting(self):
"""Get Accounting data from Mikrotik"""
# Check if accounting and account-local-traffic is enabled
accounting_enabled, local_traffic_enabled = self.api.is_accounting_and_local_traffic_enabled()
(
accounting_enabled,
local_traffic_enabled,
) = self.api.is_accounting_and_local_traffic_enabled()
uom_type, uom_div = self._get_unit_of_measurement()
# Build missing hosts from main hosts dict
for uid, vals in self.data["host"].items():
if uid not in self.data["accounting"]:
self.data["accounting"][uid] = {
'address': vals['address'],
'mac-address': vals['mac-address'],
'host-name': vals['host-name'],
'tx-rx-attr': uom_type,
'available': False,
'local_accounting': False
"address": vals["address"],
"mac-address": vals["mac-address"],
"host-name": vals["host-name"],
"tx-rx-attr": uom_type,
"available": False,
"local_accounting": False,
}
_LOGGER.debug(f"Working with {len(self.data['accounting'])} accounting devices")
# Build temp accounting values dict with ip address as key
tmp_accounting_values = {}
for uid, vals in self.data['accounting'].items():
tmp_accounting_values[vals['address']] = {
for uid, vals in self.data["accounting"].items():
tmp_accounting_values[vals["address"]] = {
"wan-tx": 0,
"wan-rx": 0,
"lan-tx": 0,
"lan-rx": 0
"lan-rx": 0,
}
time_diff = self.api.take_accounting_snapshot()
@ -1002,57 +1088,75 @@ class MikrotikControllerData:
)
for item in accounting_data.values():
source_ip = str(item.get('src-address')).strip()
destination_ip = str(item.get('dst-address')).strip()
bits_count = int(str(item.get('bytes')).strip()) * 8
source_ip = str(item.get("src-address")).strip()
destination_ip = str(item.get("dst-address")).strip()
bits_count = int(str(item.get("bytes")).strip()) * 8
if self._address_part_of_local_network(source_ip) and \
self._address_part_of_local_network(destination_ip):
if self._address_part_of_local_network(
source_ip
) and self._address_part_of_local_network(destination_ip):
# LAN TX/RX
if source_ip in tmp_accounting_values:
tmp_accounting_values[source_ip]['lan-tx'] += bits_count
tmp_accounting_values[source_ip]["lan-tx"] += bits_count
if destination_ip in tmp_accounting_values:
tmp_accounting_values[destination_ip]['lan-rx'] += bits_count
elif self._address_part_of_local_network(source_ip) and \
not self._address_part_of_local_network(destination_ip):
tmp_accounting_values[destination_ip]["lan-rx"] += bits_count
elif self._address_part_of_local_network(
source_ip
) and not self._address_part_of_local_network(destination_ip):
# WAN TX
if source_ip in tmp_accounting_values:
tmp_accounting_values[source_ip]['wan-tx'] += bits_count
elif not self._address_part_of_local_network(source_ip) and \
self._address_part_of_local_network(destination_ip) and \
destination_ip in tmp_accounting_values:
tmp_accounting_values[source_ip]["wan-tx"] += bits_count
elif (
not self._address_part_of_local_network(source_ip)
and self._address_part_of_local_network(destination_ip)
and destination_ip in tmp_accounting_values
):
# WAN RX
tmp_accounting_values[destination_ip]['wan-rx'] += bits_count
tmp_accounting_values[destination_ip]["wan-rx"] += bits_count
# Calculate real throughput and transform it to appropriate unit
# Also handle availability of accounting and local_accounting from Mikrotik
for addr, vals in tmp_accounting_values.items():
uid = self._get_accounting_uid_by_ip(addr)
if not uid:
_LOGGER.warning(f"Address {addr} not found in accounting data, skipping update")
_LOGGER.warning(
f"Address {addr} not found in accounting data, skipping update"
)
continue
self.data['accounting'][uid]['tx-rx-attr'] = uom_type
self.data['accounting'][uid]['available'] = accounting_enabled
self.data['accounting'][uid]['local_accounting'] = local_traffic_enabled
self.data["accounting"][uid]["tx-rx-attr"] = uom_type
self.data["accounting"][uid]["available"] = accounting_enabled
self.data["accounting"][uid]["local_accounting"] = local_traffic_enabled
if not accounting_enabled:
# Skip calculation for WAN and LAN if accounting is disabled
continue
self.data['accounting'][uid]['wan-tx'] = \
round(vals['wan-tx'] / time_diff * uom_div, 2) if vals['wan-tx'] else 0.0
self.data['accounting'][uid]['wan-rx'] = \
round(vals['wan-rx'] / time_diff * uom_div, 2) if vals['wan-rx'] else 0.0
self.data["accounting"][uid]["wan-tx"] = (
round(vals["wan-tx"] / time_diff * uom_div, 2)
if vals["wan-tx"]
else 0.0
)
self.data["accounting"][uid]["wan-rx"] = (
round(vals["wan-rx"] / time_diff * uom_div, 2)
if vals["wan-rx"]
else 0.0
)
if not local_traffic_enabled:
# Skip calculation for LAN if LAN accounting is disabled
continue
self.data['accounting'][uid]['lan-tx'] = \
round(vals['lan-tx'] / time_diff * uom_div, 2) if vals['lan-tx'] else 0.0
self.data['accounting'][uid]['lan-rx'] = \
round(vals['lan-rx'] / time_diff * uom_div, 2) if vals['lan-rx'] else 0.0
self.data["accounting"][uid]["lan-tx"] = (
round(vals["lan-tx"] / time_diff * uom_div, 2)
if vals["lan-tx"]
else 0.0
)
self.data["accounting"][uid]["lan-rx"] = (
round(vals["lan-rx"] / time_diff * uom_div, 2)
if vals["lan-rx"]
else 0.0
)
# ---------------------------
# _get_unit_of_measurement
@ -1088,8 +1192,8 @@ class MikrotikControllerData:
# _get_accounting_uid_by_ip
# ---------------------------
def _get_accounting_uid_by_ip(self, requested_ip):
for mac, vals in self.data['accounting'].items():
if vals.get('address') is requested_ip:
for mac, vals in self.data["accounting"].items():
if vals.get("address") is requested_ip:
return mac
return None

41
custom_components/mikrotik_router/mikrotikapi.py
View file

@ -1,11 +1,9 @@
"""Mikrotik API for Mikrotik Router."""
import importlib
import logging
import os
import ssl
import sys
import time
from time import time
from threading import Lock
from voluptuous import Optional
@ -65,7 +63,7 @@ class MikrotikAPI:
def connection_check(self) -> bool:
"""Check if mikrotik is connected"""
if not self._connected or not self._connection:
if self._connection_epoch > time.time() - self._connection_retry_sec:
if self._connection_epoch > time() - self._connection_retry_sec:
return False
if not self.connect():
@ -76,13 +74,18 @@ class MikrotikAPI:
# ---------------------------
# disconnect
# ---------------------------
def disconnect(self, location="unknown", error="unknown"):
def disconnect(self, location="unknown", error=None):
"""Disconnect from Mikrotik device."""
if not error:
error = "unknown"
if not self.connection_error_reported:
if location == "unknown":
_LOGGER.error("Mikrotik %s connection closed", self._host)
else:
_LOGGER.error("Mikrotik %s error while %s : %s", self._host, location, error)
_LOGGER.error(
"Mikrotik %s error while %s : %s", self._host, location, error
)
self.connection_error_reported = True
@ -97,7 +100,7 @@ class MikrotikAPI:
"""Connect to Mikrotik device."""
self.error = ""
self._connected = None
self._connection_epoch = time.time()
self._connection_epoch = time()
kwargs = {
"encoding": self._encoding,
@ -129,8 +132,7 @@ class MikrotikAPI:
) as api_error:
if not self.connection_error_reported:
_LOGGER.error(
"Mikrotik %s error while connecting: %s", self._host,
api_error
"Mikrotik %s error while connecting: %s", self._host, api_error
)
self.connection_error_reported = True
@ -141,8 +143,7 @@ class MikrotikAPI:
except:
if not self.connection_error_reported:
_LOGGER.error(
"Mikrotik %s error while connecting: %s", self._host,
"Unknown"
"Mikrotik %s error while connecting: %s", self._host, "Unknown"
)
self.connection_error_reported = True
@ -282,7 +283,12 @@ class MikrotikAPI:
self.lock.release()
if not entry_found:
_LOGGER.error("Mikrotik %s Update parameter %s with value %s not found", self._host, param, value)
_LOGGER.error(
"Mikrotik %s Update parameter %s with value %s not found",
self._host,
param,
value,
)
return True
@ -459,7 +465,6 @@ class MikrotikAPI:
@staticmethod
def _current_milliseconds():
from time import time
return int(round(time() * 1000))
def is_accounting_and_local_traffic_enabled(self) -> (bool, bool):
@ -475,15 +480,15 @@ class MikrotikAPI:
return False, False
for item in response:
if 'enabled' not in item:
if "enabled" not in item:
continue
if not item['enabled']:
if not item["enabled"]:
return False, False
for item in response:
if 'account-local-traffic' not in item:
if "account-local-traffic" not in item:
continue
if not item['account-local-traffic']:
if not item["account-local-traffic"]:
return True, False
return True, True
@ -502,7 +507,7 @@ class MikrotikAPI:
self.lock.acquire()
try:
# Prepare command
take = accounting('snapshot/take')
take = accounting("snapshot/take")
except librouteros.exceptions.ConnectionClosed:
self.disconnect()
self.lock.release()

39
custom_components/mikrotik_router/sensor.py
View file

@ -2,13 +2,13 @@
import logging
from homeassistant.const import (CONF_NAME, ATTR_ATTRIBUTION, ATTR_DEVICE_CLASS)
from homeassistant.const import CONF_NAME, ATTR_ATTRIBUTION, ATTR_DEVICE_CLASS
from homeassistant.core import callback
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
from .const import (DOMAIN, DATA_CLIENT, ATTRIBUTION)
from .const import DOMAIN, DATA_CLIENT, ATTRIBUTION
_LOGGER = logging.getLogger(__name__)
@ -117,11 +117,7 @@ SENSOR_TYPES = {
},
}
DEVICE_ATTRIBUTES_ACCOUNTING = [
"address",
"mac-address",
"host-name"
]
DEVICE_ATTRIBUTES_ACCOUNTING = ["address", "mac-address", "host-name"]
# ---------------------------
@ -165,8 +161,7 @@ def update_items(inst, mikrotik_controller, async_add_entities, sensors):
continue
sensors[item_id] = MikrotikControllerSensor(
mikrotik_controller=mikrotik_controller, inst=inst,
sensor=sensor
mikrotik_controller=mikrotik_controller, inst=inst, sensor=sensor
)
new_sensors.append(sensors[item_id])
@ -197,7 +192,10 @@ def update_items(inst, mikrotik_controller, async_add_entities, sensors):
sensors[item_id].async_schedule_update_ha_state()
continue
if SENSOR_TYPES[sensor][ATTR_ATTR] in mikrotik_controller.data['accounting'][uid].keys():
if (
SENSOR_TYPES[sensor][ATTR_ATTR]
in mikrotik_controller.data["accounting"][uid].keys()
):
sensors[item_id] = MikrotikAccountingSensor(
mikrotik_controller=mikrotik_controller,
inst=inst,
@ -319,8 +317,7 @@ class MikrotikControllerTrafficSensor(MikrotikControllerSensor):
"""Initialize."""
super().__init__(mikrotik_controller, inst, sensor)
self._uid = uid
self._data = mikrotik_controller.data[SENSOR_TYPES[sensor][ATTR_PATH]][
uid]
self._data = mikrotik_controller.data[SENSOR_TYPES[sensor][ATTR_PATH]][uid]
@property
def name(self):
@ -336,8 +333,7 @@ class MikrotikControllerTrafficSensor(MikrotikControllerSensor):
def device_info(self):
"""Return a port description for device registry."""
info = {
"connections": {
(CONNECTION_NETWORK_MAC, self._data["port-mac-address"])},
"connections": {(CONNECTION_NETWORK_MAC, self._data["port-mac-address"])},
"manufacturer": self._ctrl.data["resource"]["platform"],
"model": self._ctrl.data["resource"]["board-name"],
"name": f"{self._inst} {self._data['default-name']}",
@ -382,20 +378,23 @@ class MikrotikAccountingSensor(MikrotikControllerSensor):
"""Return if controller and accounting feature in Mikrotik is available.
Additional check for lan-tx/rx sensors
"""
if self._attr in ['lan-tx', 'lan-rx']:
return self._ctrl.connected() and self._data['available'] and self._data['local_accounting']
if self._attr in ["lan-tx", "lan-rx"]:
return (
self._ctrl.connected()
and self._data["available"]
and self._data["local_accounting"]
)
else:
return self._ctrl.connected() and self._data['available']
return self._ctrl.connected() and self._data["available"]
@property
def device_info(self):
"""Return a accounting description for device registry."""
info = {
"connections": {
(CONNECTION_NETWORK_MAC, self._data["mac-address"])},
"connections": {(CONNECTION_NETWORK_MAC, self._data["mac-address"])},
"manufacturer": self._ctrl.data["resource"]["platform"],
"model": self._ctrl.data["resource"]["board-name"],
"name": self._data['host-name'],
"name": self._data["host-name"],
}
return info

20
custom_components/mikrotik_router/switch.py
View file

@ -2,13 +2,13 @@
import logging
from homeassistant.components.switch import SwitchDevice
from homeassistant.const import (CONF_NAME, ATTR_ATTRIBUTION)
from homeassistant.const import CONF_NAME, ATTR_ATTRIBUTION
from homeassistant.core import callback
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.restore_state import RestoreEntity
from .const import (DOMAIN, DATA_CLIENT, ATTRIBUTION)
from .const import DOMAIN, DATA_CLIENT, ATTRIBUTION
_LOGGER = logging.getLogger(__name__)
@ -204,8 +204,7 @@ class MikrotikControllerPortSwitch(MikrotikControllerSwitch):
def device_info(self):
"""Return a port description for device registry."""
info = {
"connections": {
(CONNECTION_NETWORK_MAC, self._data["port-mac-address"])},
"connections": {(CONNECTION_NETWORK_MAC, self._data["port-mac-address"])},
"manufacturer": self._ctrl.data["resource"]["platform"],
"model": self._ctrl.data["resource"]["board-name"],
"name": f"{self._inst} {self._data['default-name']}",
@ -375,8 +374,7 @@ class MikrotikControllerScriptSwitch(MikrotikControllerSwitch):
async def async_added_to_hass(self):
"""Script switch entity created."""
_LOGGER.debug("New script switch %s (%s)", self._inst,
self._data["name"])
_LOGGER.debug("New script switch %s (%s)", self._inst, self._data["name"])
@property
def name(self) -> str:
@ -512,10 +510,7 @@ class MikrotikControllerQueueSwitch(MikrotikControllerSwitch):
param = ".id"
value = None
for uid in self._ctrl.data["queue"]:
if (
self._ctrl.data["queue"][uid]["name"]
== f"{self._data['name']}"
):
if self._ctrl.data["queue"][uid]["name"] == f"{self._data['name']}":
value = self._ctrl.data["queue"][uid][".id"]
mod_param = "disabled"
@ -529,10 +524,7 @@ class MikrotikControllerQueueSwitch(MikrotikControllerSwitch):
param = ".id"
value = None
for uid in self._ctrl.data["queue"]:
if (
self._ctrl.data["queue"][uid]["name"]
== f"{self._data['name']}"
):
if self._ctrl.data["queue"][uid]["name"] == f"{self._data['name']}":
value = self._ctrl.data["queue"][uid][".id"]
mod_param = "disabled"