diff --git a/CHANGELOG.md b/CHANGELOG.md index d986c89..186cc40 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,23 @@ # Changelog +## 2.1.0 + +Major refactor: + +- Code cleanup +- Fix thread safe issues +- Fix typos +- Improve performance +- Add util for translations +- Removed service of update configuration + +New components: + +- Consider Away Interval - Number +- Update API Interval - Number +- Update Entities Interval - Number +- Log Incoming Messages - Switch + ## 2.0.32 - Ignore interfaces that were removed diff --git a/custom_components/edgeos/__init__.py b/custom_components/edgeos/__init__.py index d4f4491..18aaf8a 100644 --- a/custom_components/edgeos/__init__.py +++ b/custom_components/edgeos/__init__.py @@ -1,60 +1,108 @@ """ -Init. +This component provides support for EdgeOS based devices. +For more details about this component, please refer to the documentation at +https://github.com/elad-bar/ha-EdgeOS """ import logging import sys from homeassistant.config_entries import ConfigEntry +from homeassistant.const import EVENT_HOMEASSISTANT_START from homeassistant.core import HomeAssistant -from .component.helpers import async_set_ha, clear_ha, get_ha -from .configuration.helpers.const import DOMAIN +from .common.consts import DEFAULT_NAME, DOMAIN +from .common.entity_descriptions import PLATFORMS +from .managers.config_manager import ConfigManager +from .managers.coordinator import Coordinator +from .managers.password_manager import PasswordManager +from .models.exceptions import LoginError _LOGGER = logging.getLogger(__name__) -async def async_setup(hass, config): +async def async_setup(_hass, _config): return True async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: - """Set up a component.""" + """Set up a EdgeOS component.""" initialized = False try: - _LOGGER.debug(f"Starting async_setup_entry of {DOMAIN}") - entry.add_update_listener(async_options_updated) + _LOGGER.debug("Setting up") + entry_config = {key: entry.data[key] for key in entry.data} - await async_set_ha(hass, entry) + _LOGGER.debug("Starting up password manager") + await PasswordManager.decrypt(hass, entry_config, entry.entry_id) - initialized = True + _LOGGER.debug("Starting up configuration manager") + config_manager = ConfigManager(hass, entry) + await config_manager.initialize(entry_config) + + is_initialized = config_manager.is_initialized + + if is_initialized: + _LOGGER.debug("Starting up coordinator") + coordinator = Coordinator(hass, config_manager) + + hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + + if hass.is_running: + _LOGGER.debug("Initializing coordinator") + await coordinator.initialize() + + else: + _LOGGER.debug("Registering listener for HA started event") + hass.bus.async_listen_once( + EVENT_HOMEASSISTANT_START, coordinator.on_home_assistant_start + ) + + _LOGGER.info("Finished loading integration") + + initialized = is_initialized + + _LOGGER.debug(f"Setup status: {is_initialized}") + + except LoginError: + _LOGGER.info(f"Failed to login {DEFAULT_NAME} API, cannot log integration") except Exception as ex: exc_type, exc_obj, tb = sys.exc_info() line_number = tb.tb_lineno - _LOGGER.error(f"Failed to load integration, error: {ex}, line: {line_number}") + _LOGGER.error( + f"Failed to load {DEFAULT_NAME}, error: {ex}, line: {line_number}" + ) return initialized async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry): """Unload a config entry.""" - ha = get_ha(hass, entry.entry_id) + _LOGGER.info(f"Unloading {DOMAIN} integration, Entry ID: {entry.entry_id}") + + coordinator: Coordinator = hass.data[DOMAIN][entry.entry_id] - if ha is not None: - await ha.async_remove(entry) + await coordinator.terminate() - clear_ha(hass, entry.entry_id) + for platform in PLATFORMS: + await hass.config_entries.async_forward_entry_unload(entry, platform) + + del hass.data[DOMAIN][entry.entry_id] return True -async def async_options_updated(hass: HomeAssistant, entry: ConfigEntry): - """Triggered by config entry options updates.""" - _LOGGER.info(f"async_options_updated, Entry: {entry.as_dict()} ") +async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry): + """Unload a config entry.""" + _LOGGER.info(f"Removing {DOMAIN} integration, Entry ID: {entry.entry_id}") + + entry_id = entry.entry_id + + coordinator: Coordinator = hass.data[DOMAIN][entry_id] + + await coordinator.config_manager.remove(entry_id) - ha = get_ha(hass, entry.entry_id) + result = await async_unload_entry(hass, entry) - if ha is not None: - await ha.async_update_entry(entry) + return result diff --git a/custom_components/edgeos/binary_sensor.py b/custom_components/edgeos/binary_sensor.py index baa2bd1..3e35170 100644 --- a/custom_components/edgeos/binary_sensor.py +++ b/custom_components/edgeos/binary_sensor.py @@ -1,34 +1,58 @@ -""" -Support for binary sensors. -""" -from __future__ import annotations - import logging -from .core.components.binary_sensor import CoreBinarySensor -from .core.helpers.setup_base_entry import async_setup_base_entry +from homeassistant.components.binary_sensor import BinarySensorEntity +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ATTR_ICON, Platform +from homeassistant.core import HomeAssistant + +from .common.base_entity import IntegrationBaseEntity, async_setup_base_entry +from .common.consts import ATTR_ATTRIBUTES, ATTR_IS_ON +from .common.entity_descriptions import IntegrationBinarySensorEntityDescription +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator _LOGGER = logging.getLogger(__name__) -async def async_setup_entry(hass, config_entry, async_add_devices): - """Set up the component.""" +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities +): await async_setup_base_entry( hass, - config_entry, - async_add_devices, - CoreBinarySensor.get_domain(), - CoreBinarySensor.get_component, + entry, + Platform.BINARY_SENSOR, + IntegrationBinarySensorEntity, + async_add_entities, ) -async def async_unload_entry(hass, config_entry): - _LOGGER.info( - f"Unload entry for {CoreBinarySensor.get_domain()} domain: {config_entry}" - ) +class IntegrationBinarySensorEntity(IntegrationBaseEntity, BinarySensorEntity): + """Representation of a sensor.""" + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationBinarySensorEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(hass, entity_description, coordinator, device_type, item_id) + + self._attr_device_class = entity_description.device_class + + def update_component(self, data): + """Fetch new state parameters for the sensor.""" + if data is not None: + is_on = data.get(ATTR_IS_ON) + attributes = data.get(ATTR_ATTRIBUTES) + icon = data.get(ATTR_ICON) - return True + self._attr_is_on = is_on + self._attr_extra_state_attributes = attributes + if icon is not None: + self._attr_icon = icon -async def async_remove_entry(hass, entry) -> None: - _LOGGER.info(f"Remove entry for {CoreBinarySensor.get_domain()} entry: {entry}") + else: + self._attr_is_on = None diff --git a/custom_components/edgeos/common/base_entity.py b/custom_components/edgeos/common/base_entity.py new file mode 100644 index 0000000..17967cd --- /dev/null +++ b/custom_components/edgeos/common/base_entity.py @@ -0,0 +1,173 @@ +import logging +import sys +from typing import Any + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import Platform +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.update_coordinator import CoordinatorEntity +from homeassistant.util import slugify + +from ..managers.coordinator import Coordinator +from .consts import ADD_COMPONENT_SIGNALS, DOMAIN +from .entity_descriptions import IntegrationEntityDescription, get_entity_descriptions +from .enums import DeviceTypes + +_LOGGER = logging.getLogger(__name__) + + +async def async_setup_base_entry( + hass: HomeAssistant, + entry: ConfigEntry, + platform: Platform, + entity_type: type, + async_add_entities, +): + @callback + def _async_handle_device( + entry_id: str, device_type: DeviceTypes, item_id: str | None = None + ): + if entry.entry_id != entry_id: + return + + try: + coordinator = hass.data[DOMAIN][entry.entry_id] + + if device_type == DeviceTypes.DEVICE: + is_monitored = coordinator.config_manager.get_monitored_device(item_id) + + elif device_type == DeviceTypes.INTERFACE: + is_monitored = coordinator.config_manager.get_monitored_interface( + item_id + ) + + else: + is_monitored = True + + entity_descriptions = get_entity_descriptions( + platform, device_type, is_monitored + ) + + entities = [ + entity_type(hass, entity_description, coordinator, device_type, item_id) + for entity_description in entity_descriptions + ] + + async_add_entities(entities, True) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to initialize {platform}, Error: {ex}, Line: {line_number}" + ) + + for add_component_signal in ADD_COMPONENT_SIGNALS: + entry.async_on_unload( + async_dispatcher_connect(hass, add_component_signal, _async_handle_device) + ) + + +class IntegrationBaseEntity(CoordinatorEntity): + _entity_description: IntegrationEntityDescription + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(coordinator) + + try: + self.hass = hass + self._item_id = item_id + self._device_type = device_type + + device_info = coordinator.get_device_info(entity_description, item_id) + + entity_name = coordinator.config_manager.get_entity_name( + entity_description, device_info + ) + + unique_id_parts = [ + DOMAIN, + entity_description.platform, + entity_description.key, + item_id, + ] + + unique_id_parts_clean = [ + unique_id_part + for unique_id_part in unique_id_parts + if unique_id_part is not None + ] + + unique_id = slugify("_".join(unique_id_parts_clean)) + + self.entity_description = entity_description + self._entity_description = entity_description + + self._attr_device_info = device_info + self._attr_name = entity_name + self._attr_unique_id = unique_id + + self._data = {} + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to initialize {entity_description}, Error: {ex}, Line: {line_number}" + ) + + @property + def _local_coordinator(self) -> Coordinator: + return self.coordinator + + @property + def data(self) -> dict | None: + return self._data + + async def async_execute_device_action(self, key: str, *kwargs: Any): + async_device_action = self._local_coordinator.get_device_action( + self._entity_description, self._item_id, key + ) + + if self._item_id is None: + await async_device_action(self._entity_description, *kwargs) + + else: + await async_device_action(self._entity_description, self._item_id, *kwargs) + + await self.coordinator.async_request_refresh() + + def update_component(self, data): + pass + + def _handle_coordinator_update(self) -> None: + """Fetch new state parameters for the sensor.""" + try: + new_data = self._local_coordinator.get_data( + self._entity_description, self._item_id + ) + + if self._data != new_data: + self.update_component(new_data) + + self._data = new_data + + self.async_write_ha_state() + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to update {self.unique_id}, Error: {ex}, Line: {line_number}" + ) diff --git a/custom_components/edgeos/common/connectivity_status.py b/custom_components/edgeos/common/connectivity_status.py new file mode 100644 index 0000000..a70eb9a --- /dev/null +++ b/custom_components/edgeos/common/connectivity_status.py @@ -0,0 +1,43 @@ +import logging + +from homeassistant.const import StrEnum + + +class ConnectivityStatus(StrEnum): + NotConnected = "Not connected" + Connecting = "Establishing connection to API" + Connected = "Connected to the API" + TemporaryConnected = "Connected with temporary API key" + Failed = "Failed to access API" + InvalidCredentials = "Invalid credentials" + MissingAPIKey = "Permanent API Key was not found" + Disconnected = "Disconnected by the system" + NotFound = "API Not found" + + @staticmethod + def get_log_level(status: StrEnum) -> int: + if status in [ + ConnectivityStatus.Connected, + ConnectivityStatus.Connecting, + ConnectivityStatus.Disconnected, + ConnectivityStatus.TemporaryConnected, + ]: + return logging.INFO + elif status in [ConnectivityStatus.NotConnected]: + return logging.WARNING + else: + return logging.ERROR + + @staticmethod + def get_ha_error(status: str) -> str | None: + errors = { + str(ConnectivityStatus.InvalidCredentials): "invalid_admin_credentials", + str(ConnectivityStatus.TemporaryConnected): "missing_permanent_api_key", + str(ConnectivityStatus.MissingAPIKey): "missing_permanent_api_key", + str(ConnectivityStatus.Failed): "invalid_server_details", + str(ConnectivityStatus.NotFound): "invalid_server_details", + } + + error_id = errors.get(status) + + return error_id diff --git a/custom_components/edgeos/component/helpers/const.py b/custom_components/edgeos/common/consts.py similarity index 81% rename from custom_components/edgeos/component/helpers/const.py rename to custom_components/edgeos/common/consts.py index e1341ec..9462c0f 100644 --- a/custom_components/edgeos/component/helpers/const.py +++ b/custom_components/edgeos/common/consts.py @@ -6,15 +6,77 @@ import aiohttp import voluptuous as vol +from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME import homeassistant.helpers.config_validation as cv -from .enums import InterfaceTypes +from .enums import DynamicInterfaceTypes + +ENTITY_CONFIG_ENTRY_ID = "entry_id" + +HA_NAME = "homeassistant" + +DOMAIN = "edgeos" +DEFAULT_NAME = "EdgeOS" +MANUFACTURER = "Ubiquiti" + +STORAGE_DATA_KEY = "key" + +SIGNAL_INTERFACE_DISCOVERED = f"{DOMAIN}_INTERFACE_DISCOVERED_SIGNAL" +SIGNAL_INTERFACE_ADDED = f"{DOMAIN}_INTERFACE_ADDED_SIGNAL" + +SIGNAL_DEVICE_DISCOVERED = f"{DOMAIN}_DEVICE_DISCOVERED_SIGNAL" +SIGNAL_DEVICE_ADDED = f"{DOMAIN}_DEVICE_ADDED_SIGNAL" + +SIGNAL_SYSTEM_DISCOVERED = f"{DOMAIN}_SYSTEM_DISCOVERED_SIGNAL" +SIGNAL_SYSTEM_ADDED = f"{DOMAIN}_SYSTEM_ADDED_SIGNAL" + +SIGNAL_DATA_CHANGED = f"{DOMAIN}_DATA_CHANGED_SIGNAL" + +SIGNAL_WS_READY = f"{DOMAIN}_WS_READY_SIGNAL" +SIGNAL_WS_STATUS = f"{DOMAIN}_WS_STATUS_SIGNAL" +SIGNAL_API_STATUS = f"{DOMAIN}_API_STATUS_SIGNAL" + +ADD_COMPONENT_SIGNALS = [ + SIGNAL_INTERFACE_ADDED, + SIGNAL_DEVICE_ADDED, + SIGNAL_SYSTEM_ADDED, +] + +DATA_KEYS = [CONF_HOST, CONF_USERNAME, CONF_PASSWORD] + +MAXIMUM_RECONNECT = 3 +CONFIGURATION_FILE = f"{DOMAIN}.config.json" + +INVALID_TOKEN_SECTION = "https://github.com/elad-bar/ha-edgeos#invalid-token" + +API_URL_TEMPLATE = "https://{}" +WEBSOCKET_URL_TEMPLATE = "wss://{}/ws/stats" + +COOKIE_PHPSESSID = "PHPSESSID" +COOKIE_BEAKER_SESSION_ID = "beaker.session.id" +COOKIE_CSRF_TOKEN = "X-CSRF-TOKEN" + +HEADER_CSRF_TOKEN = "X-Csrf-token" +EMPTY_STRING = "" +CONF_TITLE = "title" ATTR_FRIENDLY_NAME = "friendly_name" +ATTR_ATTRIBUTES = "attributes" +ATTR_ACTIONS = "actions" +ATTR_IS_ON = "is_on" +ATTR_LAST_ACTIVITY = "last activity" +ATTR_HOSTNAME = "hostname" + +ACTION_ENTITY_TURN_ON = "turn_on" +ACTION_ENTITY_TURN_OFF = "turn_off" +ACTION_ENTITY_SET_NATIVE_VALUE = "set_native_value" +ACTION_ENTITY_SELECT_OPTION = "select" CONF_DEVICE_ID = "device_id" WS_MAX_MSG_SIZE = 0 +DISCONNECT_INTERVAL = 5 + WS_RECONNECT_INTERVAL = timedelta(seconds=30) WS_TIMEOUT = timedelta(minutes=1) WS_WARNING_INTERVAL = timedelta(seconds=95) @@ -25,6 +87,8 @@ DEFAULT_UPDATE_ENTITIES_INTERVAL = timedelta(seconds=1) DEFAULT_HEARTBEAT_INTERVAL = timedelta(seconds=50) DEFAULT_CONSIDER_AWAY_INTERVAL = timedelta(minutes=3) +API_RECONNECT_INTERVAL = timedelta(seconds=30) +HEARTBEAT_INTERVAL = timedelta(seconds=25) STORAGE_DATA_FILE_CONFIG = "config" @@ -76,6 +140,7 @@ API_DATA_SYS_INFO = "sys_info" API_DATA_DHCP_LEASES = "dhcp-leases" +DATA_SYSTEM_SYSTEM = "system" DATA_SYSTEM_SERVICE = "service" DATA_SYSTEM_SERVICE_DHCP_SERVER = "dhcp-server" @@ -187,7 +252,7 @@ INTERFACE_DATA_UP = "up" INTERFACE_DATA_LINK_UP = "l1up" INTERFACE_DATA_MAC = "mac" -INTERFACE_DATA_HANDLER = "handler" +INTERFACE_DATA_IS_SUPPORTED = "is_supported" DEVICE_DATA_NAME = "hostname" DEVICE_DATA_DOMAIN = "domain" @@ -268,15 +333,13 @@ aiohttp.WSMsgType.CLOSING, ] -SPECIAL_INTERFACES = { - InterfaceTypes.PPPOE_PREFIX: "Internet Dail-Up", - InterfaceTypes.SWITCH_PREFIX: "Switch", - InterfaceTypes.VIRTUAL_TUNNEL_PREFIX: "Virtual Tunnel", - InterfaceTypes.OPEN_VPN_PREFIX: "OpenVPN", - InterfaceTypes.BONDING_PREFIX: "VLAN", -} - -IGNORED_INTERFACES = [InterfaceTypes.LOOPBACK] +INTERFACE_DYNAMIC_SUPPORTED = [ + DynamicInterfaceTypes.PPPOE, + DynamicInterfaceTypes.SWITCH, + DynamicInterfaceTypes.VIRTUAL_TUNNEL, + DynamicInterfaceTypes.OPEN_VPN, + DynamicInterfaceTypes.BONDING, +] RECEIVED_RATE_PREFIX = "Received Rate" RECEIVED_TRAFFIC_PREFIX = "Received Traffic" diff --git a/custom_components/edgeos/common/entity_descriptions.py b/custom_components/edgeos/common/entity_descriptions.py new file mode 100644 index 0000000..96978f9 --- /dev/null +++ b/custom_components/edgeos/common/entity_descriptions.py @@ -0,0 +1,312 @@ +from copy import copy +from dataclasses import dataclass +from typing import Callable + +from custom_components.edgeos.common.enums import ( + DeviceTypes, + EntityKeys, + UnitOfInterface, +) +from homeassistant.components.binary_sensor import ( + BinarySensorDeviceClass, + BinarySensorEntityDescription, +) +from homeassistant.components.number import NumberEntityDescription +from homeassistant.components.select import SelectEntityDescription +from homeassistant.components.sensor import ( + SensorDeviceClass, + SensorEntityDescription, + SensorStateClass, +) +from homeassistant.components.switch import SwitchEntityDescription +from homeassistant.const import ( + PERCENTAGE, + EntityCategory, + Platform, + UnitOfDataRate, + UnitOfInformation, + UnitOfTime, +) +from homeassistant.helpers.entity import EntityDescription + + +@dataclass(frozen=True, kw_only=True) +class IntegrationEntityDescription(EntityDescription): + platform: Platform | None = None + filter: Callable[[bool], bool] | None = lambda is_monitored: True + + +@dataclass(frozen=True, kw_only=True) +class IntegrationBinarySensorEntityDescription( + BinarySensorEntityDescription, IntegrationEntityDescription +): + platform: Platform | None = Platform.BINARY_SENSOR + on_value: str | bool | None = None + attributes: list[str] | None = None + + +@dataclass(frozen=True, kw_only=True) +class IntegrationSensorEntityDescription( + SensorEntityDescription, IntegrationEntityDescription +): + platform: Platform | None = Platform.SENSOR + + +@dataclass(frozen=True, kw_only=True) +class IntegrationSelectEntityDescription( + SelectEntityDescription, IntegrationEntityDescription +): + platform: Platform | None = Platform.SELECT + + +@dataclass(frozen=True, kw_only=True) +class IntegrationSwitchEntityDescription( + SwitchEntityDescription, IntegrationEntityDescription +): + platform: Platform | None = Platform.SWITCH + on_value: str | bool | None = None + action_name: str | None = None + + +@dataclass(frozen=True, kw_only=True) +class IntegrationDeviceTrackerEntityDescription(IntegrationEntityDescription): + platform: Platform | None = Platform.DEVICE_TRACKER + + +@dataclass(frozen=True, kw_only=True) +class IntegrationNumberEntityDescription( + NumberEntityDescription, IntegrationEntityDescription +): + platform: Platform | None = Platform.NUMBER + + +ENTITY_DESCRIPTIONS: list[IntegrationEntityDescription] = [ + IntegrationSensorEntityDescription( + key=EntityKeys.CPU_USAGE, + native_unit_of_measurement=PERCENTAGE, + state_class=SensorStateClass.MEASUREMENT, + icon="mdi:chip", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.RAM_USAGE, + native_unit_of_measurement=PERCENTAGE, + state_class=SensorStateClass.MEASUREMENT, + icon="mdi:memory", + ), + IntegrationBinarySensorEntityDescription( + key=EntityKeys.FIRMWARE, device_class=BinarySensorDeviceClass.UPDATE + ), + IntegrationSensorEntityDescription( + key=EntityKeys.LAST_RESTART, device_class=SensorDeviceClass.TIMESTAMP + ), + IntegrationSensorEntityDescription( + key=EntityKeys.UNKNOWN_DEVICES, + native_unit_of_measurement="Devices", + state_class=SensorStateClass.MEASUREMENT, + icon="mdi:help-network-outline", + ), + IntegrationSwitchEntityDescription( + key=EntityKeys.LOG_INCOMING_MESSAGES, + entity_category=EntityCategory.CONFIG, + icon="mdi:math-log", + ), + IntegrationNumberEntityDescription( + key=EntityKeys.CONSIDER_AWAY_INTERVAL, + native_max_value=600, + native_min_value=0, + native_unit_of_measurement=UnitOfTime.SECONDS, + entity_category=EntityCategory.CONFIG, + ), + IntegrationNumberEntityDescription( + key=EntityKeys.UPDATE_ENTITIES_INTERVAL, + native_max_value=600, + native_min_value=0, + native_unit_of_measurement=UnitOfTime.SECONDS, + entity_category=EntityCategory.CONFIG, + ), + IntegrationNumberEntityDescription( + key=EntityKeys.UPDATE_API_INTERVAL, + native_max_value=600, + native_min_value=0, + native_unit_of_measurement=UnitOfTime.SECONDS, + entity_category=EntityCategory.CONFIG, + ), + IntegrationBinarySensorEntityDescription( + key=EntityKeys.INTERFACE_CONNECTED, + device_class=BinarySensorDeviceClass.CONNECTIVITY, + filter=lambda is_monitored: is_monitored, + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_RECEIVED_DROPPED, + native_unit_of_measurement=UnitOfInterface.DROPPED, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:package-variant-minus", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_SENT_DROPPED, + native_unit_of_measurement=UnitOfInterface.DROPPED, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:package-variant-minus", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_RECEIVED_ERRORS, + native_unit_of_measurement=UnitOfInterface.ERRORS, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:timeline-alert", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_SENT_ERRORS, + native_unit_of_measurement=UnitOfInterface.ERRORS, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:timeline-alert", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_RECEIVED_PACKETS, + native_unit_of_measurement=UnitOfInterface.PACKETS, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:package-up", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_SENT_PACKETS, + native_unit_of_measurement=UnitOfInterface.PACKETS, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:package-up", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_RECEIVED_RATE, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:download-network-outline", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_SENT_RATE, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:upload-network-outline", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_RECEIVED_TRAFFIC, + native_unit_of_measurement=UnitOfInformation.BYTES, + state_class=SensorStateClass.TOTAL_INCREASING, + filter=lambda is_monitored: is_monitored, + icon="mdi:download-network-outline", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.INTERFACE_SENT_TRAFFIC, + native_unit_of_measurement=UnitOfInformation.BYTES, + state_class=SensorStateClass.TOTAL_INCREASING, + filter=lambda is_monitored: is_monitored, + icon="mdi:upload-network-outline", + ), + IntegrationSwitchEntityDescription( + key=EntityKeys.INTERFACE_MONITORED, + entity_category=EntityCategory.CONFIG, + icon="mdi:monitor-eye", + ), + IntegrationSwitchEntityDescription( + key=EntityKeys.INTERFACE_STATUS, icon="mdi:monitor-eye" + ), + IntegrationSensorEntityDescription( + key=EntityKeys.DEVICE_RECEIVED_RATE, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:download-network-outline", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.DEVICE_SENT_RATE, + native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND, + state_class=SensorStateClass.MEASUREMENT, + filter=lambda is_monitored: is_monitored, + icon="mdi:upload-network-outline", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.DEVICE_RECEIVED_TRAFFIC, + native_unit_of_measurement=UnitOfInformation.BYTES, + state_class=SensorStateClass.TOTAL_INCREASING, + filter=lambda is_monitored: is_monitored, + icon="mdi:download-network-outline", + ), + IntegrationSensorEntityDescription( + key=EntityKeys.DEVICE_SENT_TRAFFIC, + native_unit_of_measurement=UnitOfInformation.BYTES, + state_class=SensorStateClass.TOTAL_INCREASING, + filter=lambda is_monitored: is_monitored, + icon="mdi:upload-network-outline", + ), + IntegrationDeviceTrackerEntityDescription( + key=EntityKeys.DEVICE_TRACKER, filter=lambda is_monitored: is_monitored + ), + IntegrationSwitchEntityDescription( + key=EntityKeys.DEVICE_MONITORED, + entity_category=EntityCategory.CONFIG, + icon="mdi:monitor-eye", + ), +] + +ENTITY_DEVICE_MAPPING = { + EntityKeys.CPU_USAGE: DeviceTypes.SYSTEM, + EntityKeys.RAM_USAGE: DeviceTypes.SYSTEM, + EntityKeys.FIRMWARE: DeviceTypes.SYSTEM, + EntityKeys.LAST_RESTART: DeviceTypes.SYSTEM, + EntityKeys.UNKNOWN_DEVICES: DeviceTypes.SYSTEM, + EntityKeys.LOG_INCOMING_MESSAGES: DeviceTypes.SYSTEM, + EntityKeys.CONSIDER_AWAY_INTERVAL: DeviceTypes.SYSTEM, + EntityKeys.UPDATE_ENTITIES_INTERVAL: DeviceTypes.SYSTEM, + EntityKeys.UPDATE_API_INTERVAL: DeviceTypes.SYSTEM, + EntityKeys.INTERFACE_CONNECTED: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_RECEIVED_DROPPED: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_SENT_DROPPED: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_RECEIVED_ERRORS: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_SENT_ERRORS: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_RECEIVED_PACKETS: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_SENT_PACKETS: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_RECEIVED_RATE: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_SENT_RATE: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_RECEIVED_TRAFFIC: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_SENT_TRAFFIC: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_MONITORED: DeviceTypes.INTERFACE, + EntityKeys.INTERFACE_STATUS: DeviceTypes.INTERFACE, + EntityKeys.DEVICE_RECEIVED_RATE: DeviceTypes.DEVICE, + EntityKeys.DEVICE_SENT_RATE: DeviceTypes.DEVICE, + EntityKeys.DEVICE_RECEIVED_TRAFFIC: DeviceTypes.DEVICE, + EntityKeys.DEVICE_SENT_TRAFFIC: DeviceTypes.DEVICE, + EntityKeys.DEVICE_TRACKER: DeviceTypes.DEVICE, + EntityKeys.DEVICE_MONITORED: DeviceTypes.DEVICE, +} + + +def get_entity_descriptions( + platform: Platform, device_type: DeviceTypes, is_monitored: bool +) -> list[IntegrationEntityDescription]: + entity_descriptions = copy(ENTITY_DESCRIPTIONS) + + result = [ + entity_description + for entity_description in entity_descriptions + if entity_description.platform == platform + and ENTITY_DEVICE_MAPPING.get(entity_description.key) == device_type + and entity_description.filter(is_monitored) + ] + + return result + + +def get_platforms() -> list[str]: + platforms = { + entity_description.platform: None for entity_description in ENTITY_DESCRIPTIONS + } + result = list(platforms.keys()) + + return result + + +PLATFORMS = get_platforms() diff --git a/custom_components/edgeos/common/enums.py b/custom_components/edgeos/common/enums.py new file mode 100644 index 0000000..b422cd4 --- /dev/null +++ b/custom_components/edgeos/common/enums.py @@ -0,0 +1,64 @@ +from enum import StrEnum + + +class DeviceTypes(StrEnum): + SYSTEM = "System" + DEVICE = "Device" + INTERFACE = "Interface" + + +class InterfaceTypes(StrEnum): + BRIDGE = "bridge" + LOOPBACK = "loopback" + ETHERNET = "ethernet" + DYNAMIC = "dynamic" + + +class DynamicInterfaceTypes(StrEnum): + PPPOE = "pppoe" + SWITCH = "switch" + VIRTUAL_TUNNEL = "vtun" + OPEN_VPN = "openvpn" + BONDING = "bond" + INTERMEDIATE_QUEUEING_DEVICE = "imq" + NETWORK_PROGRAMMING_INTERFACE = "npi" + LOOPBACK = "lo" + + +class EntityKeys(StrEnum): + CPU_USAGE = "cpu_usage" + RAM_USAGE = "ram_usage" + FIRMWARE = "firmware" + LAST_RESTART = "last_restart" + UNKNOWN_DEVICES = "unknown_devices" + LOG_INCOMING_MESSAGES = "log_incoming_messages" + CONSIDER_AWAY_INTERVAL = "consider_away_interval" + UPDATE_ENTITIES_INTERVAL = "update_entities_interval" + UPDATE_API_INTERVAL = "update_api_interval" + + INTERFACE_CONNECTED = "interface_connected" + INTERFACE_RECEIVED_DROPPED = "interface_received_dropped" + INTERFACE_SENT_DROPPED = "interface_sent_dropped" + INTERFACE_RECEIVED_ERRORS = "interface_received_errors" + INTERFACE_SENT_ERRORS = "interface_sent_errors" + INTERFACE_RECEIVED_PACKETS = "interface_received_packets" + INTERFACE_SENT_PACKETS = "interface_sent_packets" + INTERFACE_RECEIVED_RATE = "interface_received_rate" + INTERFACE_SENT_RATE = "interface_sent_rate" + INTERFACE_RECEIVED_TRAFFIC = "interface_received_traffic" + INTERFACE_SENT_TRAFFIC = "interface_sent_traffic" + INTERFACE_MONITORED = "interface_monitored" + INTERFACE_STATUS = "interface_status" + + DEVICE_RECEIVED_RATE = "device_received_rate" + DEVICE_SENT_RATE = "device_sent_rate" + DEVICE_RECEIVED_TRAFFIC = "device_received_traffic" + DEVICE_SENT_TRAFFIC = "device_sent_traffic" + DEVICE_TRACKER = "device_tracker" + DEVICE_MONITORED = "device_monitored" + + +class UnitOfInterface(StrEnum): + ERRORS = "Errors" + DROPPED = "Dropped" + PACKETS = "Packets" diff --git a/custom_components/edgeos/component/helpers/exceptions.py b/custom_components/edgeos/common/exceptions.py similarity index 70% rename from custom_components/edgeos/component/helpers/exceptions.py rename to custom_components/edgeos/common/exceptions.py index 45ee69a..54a4d74 100644 --- a/custom_components/edgeos/component/helpers/exceptions.py +++ b/custom_components/edgeos/common/exceptions.py @@ -21,3 +21,17 @@ def __init__(self, endpoint: str, status: ConnectivityStatus): self.endpoint = endpoint self.status = status + + +class AlreadyExistsError(HomeAssistantError): + title: str + + def __init__(self, title: str): + self.title = title + + +class LoginError(HomeAssistantError): + errors: dict + + def __init__(self, errors): + self.errors = errors diff --git a/custom_components/edgeos/component/__init__.py b/custom_components/edgeos/component/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/component/api/__init__.py b/custom_components/edgeos/component/api/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/component/api/storage_api.py b/custom_components/edgeos/component/api/storage_api.py deleted file mode 100644 index 4cfb694..0000000 --- a/custom_components/edgeos/component/api/storage_api.py +++ /dev/null @@ -1,191 +0,0 @@ -"""Storage handlers.""" -from __future__ import annotations - -from collections.abc import Awaitable, Callable -import logging - -from homeassistant.core import HomeAssistant -from homeassistant.helpers.json import JSONEncoder -from homeassistant.helpers.storage import Store - -from ...configuration.models.config_data import ConfigData -from ...core.api.base_api import BaseAPI -from ...core.helpers.const import DOMAIN, STORAGE_VERSION -from ...core.helpers.enums import ConnectivityStatus -from ..helpers.const import ( - DEFAULT_CONSIDER_AWAY_INTERVAL, - DEFAULT_UPDATE_API_INTERVAL, - DEFAULT_UPDATE_ENTITIES_INTERVAL, - STORAGE_DATA_CONSIDER_AWAY_INTERVAL, - STORAGE_DATA_FILE_CONFIG, - STORAGE_DATA_FILES, - STORAGE_DATA_LOG_INCOMING_MESSAGES, - STORAGE_DATA_MONITORED_DEVICES, - STORAGE_DATA_MONITORED_INTERFACES, - STORAGE_DATA_UPDATE_API_INTERVAL, - STORAGE_DATA_UPDATE_ENTITIES_INTERVAL, -) - -_LOGGER = logging.getLogger(__name__) - - -class StorageAPI(BaseAPI): - _stores: dict[str, Store] | None - _config_data: ConfigData | None - _data: dict - - def __init__( - self, - hass: HomeAssistant | None, - async_on_data_changed: Callable[[], Awaitable[None]] | None = None, - async_on_status_changed: Callable[[ConnectivityStatus], Awaitable[None]] - | None = None, - ): - super().__init__(hass, async_on_data_changed, async_on_status_changed) - - self._config_data = None - self._stores = None - self._data = {} - - @property - def _storage_config(self) -> Store: - storage = self._stores.get(STORAGE_DATA_FILE_CONFIG) - - return storage - - @property - def monitored_interfaces(self): - result = self.data.get(STORAGE_DATA_MONITORED_INTERFACES, {}) - - return result - - @property - def monitored_devices(self): - result = self.data.get(STORAGE_DATA_MONITORED_DEVICES, {}) - - return result - - @property - def log_incoming_messages(self): - result = self.data.get(STORAGE_DATA_LOG_INCOMING_MESSAGES, False) - - return result - - @property - def consider_away_interval(self): - result = self.data.get( - STORAGE_DATA_CONSIDER_AWAY_INTERVAL, - DEFAULT_CONSIDER_AWAY_INTERVAL.total_seconds(), - ) - - return result - - @property - def update_entities_interval(self): - result = self.data.get( - STORAGE_DATA_UPDATE_ENTITIES_INTERVAL, - DEFAULT_UPDATE_ENTITIES_INTERVAL.total_seconds(), - ) - - return result - - @property - def update_api_interval(self): - result = self.data.get( - STORAGE_DATA_UPDATE_API_INTERVAL, - DEFAULT_UPDATE_API_INTERVAL.total_seconds(), - ) - - return result - - async def initialize(self, config_data: ConfigData): - self._config_data = config_data - - self._initialize_storages() - - await self._async_load_configuration() - - def _initialize_storages(self): - stores = {} - - entry_id = self._config_data.entry.entry_id - - for storage_data_file in STORAGE_DATA_FILES: - file_name = f"{DOMAIN}.{entry_id}.{storage_data_file}.json" - - stores[storage_data_file] = Store( - self.hass, STORAGE_VERSION, file_name, encoder=JSONEncoder - ) - - self._stores = stores - - async def _async_load_configuration(self): - """Load the retained data from store and return de-serialized data.""" - self.data = await self._storage_config.async_load() - - if self.data is None: - self.data = { - STORAGE_DATA_MONITORED_INTERFACES: {}, - STORAGE_DATA_MONITORED_DEVICES: {}, - STORAGE_DATA_LOG_INCOMING_MESSAGES: False, - STORAGE_DATA_CONSIDER_AWAY_INTERVAL: DEFAULT_CONSIDER_AWAY_INTERVAL.total_seconds(), - STORAGE_DATA_UPDATE_ENTITIES_INTERVAL: DEFAULT_UPDATE_ENTITIES_INTERVAL.total_seconds(), - STORAGE_DATA_UPDATE_API_INTERVAL: DEFAULT_UPDATE_API_INTERVAL.total_seconds(), - } - - await self._async_save() - - _LOGGER.debug(f"Loaded configuration data: {self.data}") - - await self.set_status(ConnectivityStatus.Connected) - await self.fire_data_changed_event() - - async def _async_save(self): - """Generate dynamic data to store and save it to the filesystem.""" - _LOGGER.info(f"Save configuration, Data: {self.data}") - - await self._storage_config.async_save(self.data) - - await self.fire_data_changed_event() - - async def set_monitored_interface(self, interface_name: str, is_enabled: bool): - _LOGGER.debug(f"Set monitored interface {interface_name} to {is_enabled}") - - self.data[STORAGE_DATA_MONITORED_INTERFACES][interface_name] = is_enabled - - await self._async_save() - - async def set_monitored_device(self, device_name: str, is_enabled: bool): - _LOGGER.debug(f"Set monitored interface {device_name} to {is_enabled}") - - self.data[STORAGE_DATA_MONITORED_DEVICES][device_name] = is_enabled - - await self._async_save() - - async def set_log_incoming_messages(self, enabled: bool): - _LOGGER.debug(f"Set log incoming messages to {enabled}") - - self.data[STORAGE_DATA_LOG_INCOMING_MESSAGES] = enabled - - await self._async_save() - - async def set_consider_away_interval(self, interval: int): - _LOGGER.debug(f"Changing {STORAGE_DATA_CONSIDER_AWAY_INTERVAL}: {interval}") - - self.data[STORAGE_DATA_CONSIDER_AWAY_INTERVAL] = interval - - await self._async_save() - - async def set_update_entities_interval(self, interval: int): - _LOGGER.debug(f"Changing {STORAGE_DATA_UPDATE_ENTITIES_INTERVAL}: {interval}") - - self.data[STORAGE_DATA_UPDATE_ENTITIES_INTERVAL] = interval - - await self._async_save() - - async def set_update_api_interval(self, interval: int): - _LOGGER.debug(f"Changing {STORAGE_DATA_UPDATE_API_INTERVAL}: {interval}") - - self.data[STORAGE_DATA_UPDATE_API_INTERVAL] = interval - - await self._async_save() diff --git a/custom_components/edgeos/component/helpers/__init__.py b/custom_components/edgeos/component/helpers/__init__.py deleted file mode 100644 index 0a2935a..0000000 --- a/custom_components/edgeos/component/helpers/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -import logging -import sys - -from homeassistant.config_entries import ConfigEntry -from homeassistant.const import EVENT_HOMEASSISTANT_STOP -from homeassistant.core import Event, HomeAssistant - -from ...component.managers.home_assistant import EdgeOSHomeAssistantManager -from ...core.helpers.const import DATA - -_LOGGER = logging.getLogger(__name__) - - -async def async_set_ha(hass: HomeAssistant, entry: ConfigEntry): - try: - if DATA not in hass.data: - hass.data[DATA] = {} - - instance = EdgeOSHomeAssistantManager(hass) - - await instance.async_init(entry) - - hass.data[DATA][entry.entry_id] = instance - - async def _async_unload(_: Event) -> None: - await instance.async_unload() - - entry.async_on_unload( - hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_unload) - ) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error(f"Failed to async_set_ha, error: {ex}, line: {line_number}") - - -def get_ha(hass: HomeAssistant, entry_id) -> EdgeOSHomeAssistantManager: - ha_data = hass.data.get(DATA, {}) - ha = ha_data.get(entry_id) - - return ha - - -def clear_ha(hass: HomeAssistant, entry_id): - if DATA not in hass.data: - hass.data[DATA] = {} - - del hass.data[DATA][entry_id] diff --git a/custom_components/edgeos/component/helpers/enums.py b/custom_components/edgeos/component/helpers/enums.py deleted file mode 100644 index 7f13c37..0000000 --- a/custom_components/edgeos/component/helpers/enums.py +++ /dev/null @@ -1,24 +0,0 @@ -from enum import Enum - -from homeassistant.backports.enum import StrEnum - - -class InterfaceTypes(StrEnum): - BRIDGE = "bridge" - LOOPBACK = "loopback" - ETHERNET = "ethernet" - - PPPOE_PREFIX = "pppoe" - SWITCH_PREFIX = "switch" - VIRTUAL_TUNNEL_PREFIX = "vtun" - OPEN_VPN_PREFIX = "openvpn" - BONDING_PREFIX = "bond" - INTERMEDIATE_QUEUEING_DEVICE_PREFIX = "imq" - NETWORK_PROGRAMMING_INTERFACE_PREFIX = "npi" - LOOPBACK_PREFIX = "lo" - - -class InterfaceHandlers(Enum): - REGULAR = 0 - SPECIAL = 1 - IGNORED = 99 diff --git a/custom_components/edgeos/component/managers/__init__.py b/custom_components/edgeos/component/managers/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/component/managers/home_assistant.py b/custom_components/edgeos/component/managers/home_assistant.py deleted file mode 100644 index bcd8b57..0000000 --- a/custom_components/edgeos/component/managers/home_assistant.py +++ /dev/null @@ -1,1516 +0,0 @@ -""" -Support for HA manager. -""" -from __future__ import annotations - -from asyncio import sleep -from collections.abc import Awaitable, Callable -from datetime import datetime, timedelta -import logging -import sys - -from homeassistant.components.binary_sensor import ( - BinarySensorDeviceClass, - BinarySensorEntityDescription, -) -from homeassistant.components.homeassistant import SERVICE_RELOAD_CONFIG_ENTRY -from homeassistant.components.sensor import ( - SensorDeviceClass, - SensorEntityDescription, - SensorStateClass, -) -from homeassistant.components.switch import SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry -from homeassistant.const import ( - ATTR_FRIENDLY_NAME, - STATE_OFF, - STATE_ON, - UnitOfDataRate, - UnitOfInformation, -) -from homeassistant.core import HomeAssistant -from homeassistant.helpers.device_registry import async_get as async_get_device_registry -from homeassistant.helpers.entity import EntityCategory, EntityDescription - -from ...configuration.helpers.const import DEFAULT_NAME, DOMAIN, MANUFACTURER -from ...configuration.managers.configuration_manager import ConfigurationManager -from ...configuration.models.config_data import ConfigData -from ...core.helpers.const import ( - ACTION_CORE_ENTITY_TURN_OFF, - ACTION_CORE_ENTITY_TURN_ON, - DOMAIN_BINARY_SENSOR, - DOMAIN_DEVICE_TRACKER, - DOMAIN_SENSOR, - DOMAIN_SWITCH, - ENTITY_CONFIG_ENTRY_ID, - ENTITY_UNIQUE_ID, - HA_NAME, -) -from ...core.helpers.enums import ConnectivityStatus -from ...core.managers.home_assistant import HomeAssistantManager -from ...core.models.entity_data import EntityData -from ..api.api import IntegrationAPI -from ..api.storage_api import StorageAPI -from ..api.websocket import IntegrationWS -from ..helpers.const import ( - ADDRESS_LIST, - API_DATA_DHCP_LEASES, - API_DATA_DHCP_STATS, - API_DATA_INTERFACES, - API_DATA_SYS_INFO, - API_DATA_SYSTEM, - CONF_DEVICE_ID, - DATA_SYSTEM_SERVICE, - DATA_SYSTEM_SERVICE_DHCP_SERVER, - DEFAULT_HEARTBEAT_INTERVAL, - DEFAULT_UPDATE_API_INTERVAL, - DEVICE_DATA_MAC, - DEVICE_LIST, - DHCP_SERVER_IP_ADDRESS, - DHCP_SERVER_LEASED, - DHCP_SERVER_LEASES, - DHCP_SERVER_LEASES_CLIENT_HOSTNAME, - DHCP_SERVER_MAC_ADDRESS, - DHCP_SERVER_SHARED_NETWORK_NAME, - DHCP_SERVER_STATIC_MAPPING, - DHCP_SERVER_STATS, - DHCP_SERVER_SUBNET, - DISCOVER_DATA_FW_VERSION, - DISCOVER_DATA_PRODUCT, - FALSE_STR, - FW_LATEST_STATE_CAN_UPGRADE, - INTERFACE_DATA_ADDRESS, - INTERFACE_DATA_AGING, - INTERFACE_DATA_BRIDGE_GROUP, - INTERFACE_DATA_BRIDGED_CONNTRACK, - INTERFACE_DATA_DESCRIPTION, - INTERFACE_DATA_DUPLEX, - INTERFACE_DATA_HELLO_TIME, - INTERFACE_DATA_LINK_UP, - INTERFACE_DATA_MAC, - INTERFACE_DATA_MAX_AGE, - INTERFACE_DATA_MULTICAST, - INTERFACE_DATA_PRIORITY, - INTERFACE_DATA_PROMISCUOUS, - INTERFACE_DATA_SPEED, - INTERFACE_DATA_STP, - INTERFACE_DATA_UP, - MESSAGES_COUNTER_SECTION, - SERVICE_SCHEMA_UPDATE_CONFIGURATION, - SERVICE_UPDATE_CONFIGURATION, - STATS_DATA_RATE, - STATS_DATA_SIZE, - STATS_ICONS, - STATS_UNITS, - STORAGE_DATA_CONSIDER_AWAY_INTERVAL, - STORAGE_DATA_LOG_INCOMING_MESSAGES, - STORAGE_DATA_UPDATE_API_INTERVAL, - STORAGE_DATA_UPDATE_ENTITIES_INTERVAL, - STRING_DASH, - STRING_UNDERSCORE, - SYSTEM_DATA_DOMAIN_NAME, - SYSTEM_DATA_HOSTNAME, - SYSTEM_DATA_LOGIN, - SYSTEM_DATA_LOGIN_USER, - SYSTEM_DATA_LOGIN_USER_LEVEL, - SYSTEM_DATA_NTP, - SYSTEM_DATA_NTP_SERVER, - SYSTEM_DATA_OFFLOAD, - SYSTEM_DATA_OFFLOAD_HW_NAT, - SYSTEM_DATA_OFFLOAD_IPSEC, - SYSTEM_DATA_TIME_ZONE, - SYSTEM_DATA_TRAFFIC_ANALYSIS, - SYSTEM_DATA_TRAFFIC_ANALYSIS_DPI, - SYSTEM_DATA_TRAFFIC_ANALYSIS_EXPORT, - SYSTEM_INFO_DATA_FW_LATEST, - SYSTEM_INFO_DATA_FW_LATEST_STATE, - SYSTEM_INFO_DATA_FW_LATEST_URL, - SYSTEM_INFO_DATA_FW_LATEST_VERSION, - SYSTEM_INFO_DATA_SW_VER, - SYSTEM_STATS_DATA_CPU, - SYSTEM_STATS_DATA_MEM, - SYSTEM_STATS_DATA_UPTIME, - TRAFFIC_DATA_DEVICE_ITEMS, - TRAFFIC_DATA_DROPPED, - TRAFFIC_DATA_ERRORS, - TRAFFIC_DATA_INTERFACE_ITEMS, - TRAFFIC_DATA_PACKETS, - TRUE_STR, - USER_LEVEL_ADMIN, - WS_DISCOVER_KEY, - WS_EXPORT_KEY, - WS_INTERFACES_KEY, - WS_MESSAGES, - WS_RECONNECT_INTERVAL, - WS_SYSTEM_STATS_KEY, -) -from ..helpers.enums import InterfaceHandlers -from ..models.edge_os_device_data import EdgeOSDeviceData -from ..models.edge_os_interface_data import EdgeOSInterfaceData -from ..models.edge_os_system_data import EdgeOSSystemData - -_LOGGER = logging.getLogger(__name__) - - -class EdgeOSHomeAssistantManager(HomeAssistantManager): - def __init__(self, hass: HomeAssistant): - super().__init__(hass, DEFAULT_UPDATE_API_INTERVAL, DEFAULT_HEARTBEAT_INTERVAL) - - self._storage_api: StorageAPI = StorageAPI(self._hass) - self._api: IntegrationAPI = IntegrationAPI( - self._hass, self._api_data_changed, self._api_status_changed - ) - self._ws: IntegrationWS = IntegrationWS( - self._hass, self._ws_data_changed, self._ws_status_changed - ) - self._config_manager: ConfigurationManager | None = None - self._system: EdgeOSSystemData | None = None - self._devices: dict[str, EdgeOSDeviceData] = {} - self._devices_ip_mapping: dict[str, str] = {} - self._interfaces: dict[str, EdgeOSInterfaceData] = {} - self._can_load_components: bool = False - self._unique_messages: list[str] = [] - - @property - def hass(self) -> HomeAssistant: - return self._hass - - @property - def api(self) -> IntegrationAPI: - return self._api - - @property - def ws(self) -> IntegrationWS: - return self._ws - - @property - def storage_api(self) -> StorageAPI: - return self._storage_api - - @property - def config_data(self) -> ConfigData: - return self._config_manager.get(self.entry_id) - - @property - def system_name(self): - name = self.entry_title - - if self._system is not None and self._system.hostname is not None: - name = self._system.hostname.upper() - - return name - - async def async_send_heartbeat(self): - """Must be implemented to be able to send heartbeat to API""" - await self.ws.async_send_heartbeat() - - async def _api_data_changed(self): - if self.api.status == ConnectivityStatus.Connected: - await self._extract_api_data() - - async def _ws_data_changed(self): - if self.ws.status == ConnectivityStatus.Connected: - await self._extract_ws_data() - - async def _api_status_changed(self, status: ConnectivityStatus): - _LOGGER.info( - f"API Status changed to {status.name}, WS Status: {self.ws.status.name}" - ) - if status == ConnectivityStatus.Connected: - await self.api.async_update() - - if self.ws.status == ConnectivityStatus.NotConnected: - log_incoming_messages = self.storage_api.log_incoming_messages - await self.ws.update_api_data(self.api.data, log_incoming_messages) - - await self.ws.initialize(self.config_data) - - if status == ConnectivityStatus.Disconnected: - if self.ws.status == ConnectivityStatus.Connected: - await self.ws.terminate() - - async def _ws_status_changed(self, status: ConnectivityStatus): - _LOGGER.info( - f"WS Status changed to {status.name}, API Status: {self.api.status.name}" - ) - - api_connected = self.api.status == ConnectivityStatus.Connected - ws_connected = status == ConnectivityStatus.Connected - ws_reconnect = status in [ - ConnectivityStatus.NotConnected, - ConnectivityStatus.Failed, - ] - - self._can_load_components = ws_connected - - if ws_reconnect and api_connected: - await sleep(WS_RECONNECT_INTERVAL.total_seconds()) - - await self.ws.initialize() - - async def async_component_initialize(self, entry: ConfigEntry): - try: - self._config_manager = ConfigurationManager(self._hass, self.api) - await self._config_manager.load(entry) - - await self.storage_api.initialize(self.config_data) - - update_entities_interval = timedelta( - seconds=self.storage_api.update_entities_interval - ) - update_api_interval = timedelta( - seconds=self.storage_api.update_api_interval - ) - - _LOGGER.info( - f"Setting intervals, API: {update_api_interval}, Entities: {update_entities_interval}" - ) - self.update_intervals(update_entities_interval, update_api_interval) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to async_component_initialize, error: {ex}, line: {line_number}" - ) - - async def async_initialize_data_providers(self): - await self.storage_api.initialize(self.config_data) - - updated = False - - if self._entry is not None: - migration_data = {} - entry_options = self._entry.options - - if entry_options is not None: - for option_key in entry_options: - migration_data[option_key] = entry_options.get(option_key) - - updated = await self._update_configuration_data(migration_data) - - if updated: - _LOGGER.info("Starting configuration migration") - - data = {} - for key in self._entry.data.keys(): - value = self._entry.data.get(key) - data[key] = value - - options = {} - - self._hass.config_entries.async_update_entry( - self._entry, data=data, options=options - ) - - _LOGGER.info("Configuration migration completed, reloading integration") - - await self._reload_integration() - - else: - await self.api.initialize(self.config_data) - - async def async_stop_data_providers(self): - await self.api.terminate() - - async def async_update_data_providers(self): - try: - await self.api.async_update() - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to async_update_data_providers, Error: {ex}, Line: {line_number}" - ) - - def register_services(self, entry: ConfigEntry | None = None): - self._hass.services.async_register( - DOMAIN, - SERVICE_UPDATE_CONFIGURATION, - self._update_configuration, - SERVICE_SCHEMA_UPDATE_CONFIGURATION, - ) - - def load_devices(self): - if not self._can_load_components: - return - - self._load_main_device() - - for unique_id in self._devices: - device_item = self._get_device(unique_id) - self._load_device_device(device_item) - - for unique_id in self._interfaces: - interface_item = self._interfaces.get(unique_id) - self._load_interface_device(interface_item) - - def load_entities(self): - _LOGGER.debug("Loading entities") - - if not self._can_load_components: - return - - is_admin = self._system.user_level == USER_LEVEL_ADMIN - - self._load_unknown_devices_sensor() - self._load_cpu_sensor() - self._load_ram_sensor() - self._load_uptime_sensor() - self._load_firmware_upgrade_binary_sensor() - self._load_log_incoming_messages_switch() - - for unique_id in self._devices: - device_item = self._get_device(unique_id) - - if device_item.is_leased: - continue - - self._load_device_monitor_switch(device_item) - self._load_device_tracker(device_item) - - stats_data = device_item.get_stats() - - for stats_data_key in stats_data: - stats_data_item = stats_data.get(stats_data_key) - device_name = self.get_device_name(device_item) - - self._load_stats_sensor( - device_item.unique_id, - device_name, - stats_data_key, - stats_data_item, - self.storage_api.monitored_devices, - ) - - for unique_id in self._interfaces: - interface_item = self._interfaces.get(unique_id) - - if interface_item.handler == InterfaceHandlers.IGNORED: - continue - - if is_admin and interface_item.handler == InterfaceHandlers.REGULAR: - self._load_interface_status_switch(interface_item) - - else: - self._load_interface_status_binary_sensor(interface_item) - - self._load_interface_monitor_switch(interface_item) - self._load_interface_connected_binary_sensor(interface_item) - - stats_data = interface_item.get_stats() - - for stats_data_key in stats_data: - stats_data_item = stats_data.get(stats_data_key) - interface_name = self.get_interface_name(interface_item) - - self._load_stats_sensor( - interface_item.unique_id, - interface_name, - stats_data_key, - stats_data_item, - self.storage_api.monitored_interfaces, - ) - - def get_device_name(self, device: EdgeOSDeviceData): - return f"{self.system_name} Device {device.hostname}" - - def get_interface_name(self, interface: EdgeOSInterfaceData): - return f"{self.system_name} Interface {interface.name.upper()}" - - async def _extract_ws_data(self): - try: - interfaces_data = self.ws.data.get(WS_INTERFACES_KEY, {}) - device_data = self.ws.data.get(WS_EXPORT_KEY, {}) - - system_stats_data = self.ws.data.get(WS_SYSTEM_STATS_KEY, {}) - discovery_data = self.ws.data.get(WS_DISCOVER_KEY, {}) - - self._update_system_stats(system_stats_data, discovery_data) - - for device_ip in device_data: - device_item = self._get_device_by_ip(device_ip) - stats = device_data.get(device_ip) - - if device_item is not None: - self._update_device_stats(device_item, stats) - - for name in interfaces_data: - interface_item = self._interfaces.get(name) - stats = interfaces_data.get(name) - - if interface_item is None: - interface_data = interfaces_data.get(name) - interface_item = self._extract_interface(name, interface_data) - - self._update_interface_stats(interface_item, stats) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract WS data, Error: {ex}, Line: {line_number}" - ) - - async def _extract_api_data(self): - try: - _LOGGER.debug("Extracting API Data") - - data = self.api.data.get(API_DATA_SYSTEM, {}) - system_info = self.api.data.get(API_DATA_SYS_INFO, {}) - - self._extract_system(data, system_info) - - self._extract_unknown_devices() - - self._extract_interfaces(data) - self._extract_devices(data) - - warning_messages = [] - - if not self._system.deep_packet_inspection: - warning_messages.append("DPI (deep packet inspection) is turned off") - - if not self._system.traffic_analysis_export: - warning_messages.append("Traffic Analysis Export is turned off") - - if len(warning_messages) > 0: - warning_message = " and ".join(warning_messages) - - _LOGGER.warning( - f"Integration will not work correctly since {warning_message}" - ) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract API data, Error: {ex}, Line: {line_number}" - ) - - def get_debug_data(self) -> dict: - messages = {} - - for key in WS_MESSAGES: - message_counter = self._ws.data.get(key, 0) - counter_name = key.replace(f"-{MESSAGES_COUNTER_SECTION.lower()}", "") - - messages[counter_name] = message_counter - - data = { - API_DATA_SYSTEM: self._system, - DEVICE_LIST: self._devices, - API_DATA_INTERFACES: self._interfaces, - MESSAGES_COUNTER_SECTION: messages, - } - - return data - - def _extract_system(self, data: dict, system_info: dict): - try: - system_details = data.get(API_DATA_SYSTEM, {}) - - system_data = EdgeOSSystemData() if self._system is None else self._system - - system_data.hostname = system_details.get(SYSTEM_DATA_HOSTNAME) - system_data.timezone = system_details.get(SYSTEM_DATA_TIME_ZONE) - - ntp: dict = system_details.get(SYSTEM_DATA_NTP, {}) - system_data.ntp_servers = ntp.get(SYSTEM_DATA_NTP_SERVER) - - offload: dict = system_details.get(SYSTEM_DATA_OFFLOAD, {}) - hardware_offload = EdgeOSSystemData.is_enabled( - offload, SYSTEM_DATA_OFFLOAD_HW_NAT - ) - ipsec_offload = EdgeOSSystemData.is_enabled( - offload, SYSTEM_DATA_OFFLOAD_IPSEC - ) - - system_data.hardware_offload = hardware_offload - system_data.ipsec_offload = ipsec_offload - - traffic_analysis: dict = system_details.get( - SYSTEM_DATA_TRAFFIC_ANALYSIS, {} - ) - dpi = EdgeOSSystemData.is_enabled( - traffic_analysis, SYSTEM_DATA_TRAFFIC_ANALYSIS_DPI - ) - traffic_analysis_export = EdgeOSSystemData.is_enabled( - traffic_analysis, SYSTEM_DATA_TRAFFIC_ANALYSIS_EXPORT - ) - - system_data.deep_packet_inspection = dpi - system_data.traffic_analysis_export = traffic_analysis_export - - sw_latest = system_info.get(SYSTEM_INFO_DATA_SW_VER) - fw_latest = system_info.get(SYSTEM_INFO_DATA_FW_LATEST, {}) - - fw_latest_state = fw_latest.get(SYSTEM_INFO_DATA_FW_LATEST_STATE) - fw_latest_version = fw_latest.get(SYSTEM_INFO_DATA_FW_LATEST_VERSION) - fw_latest_url = fw_latest.get(SYSTEM_INFO_DATA_FW_LATEST_URL) - - system_data.upgrade_available = ( - fw_latest_state == FW_LATEST_STATE_CAN_UPGRADE - ) - system_data.upgrade_url = fw_latest_url - system_data.upgrade_version = fw_latest_version - - system_data.sw_version = sw_latest - - login_details = system_details.get(SYSTEM_DATA_LOGIN, {}) - users = login_details.get(SYSTEM_DATA_LOGIN_USER, {}) - current_user = users.get(self.config_data.username, {}) - system_data.user_level = current_user.get(SYSTEM_DATA_LOGIN_USER_LEVEL) - - self._system = system_data - - message = ( - f"User {self.config_data.username} level is {self._system.user_level}, " - f"Interface status switch will not be created as it requires admin role" - ) - - self.unique_log(logging.INFO, message) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract System data, Error: {ex}, Line: {line_number}" - ) - - def _extract_interfaces(self, data: dict): - try: - interface_types = data.get(API_DATA_INTERFACES, {}) - - for interface_type in interface_types: - interfaces = interface_types.get(interface_type) - - if interfaces is not None: - for interface_name in interfaces: - interface_data = interfaces.get(interface_name, {}) - self._extract_interface( - interface_name, interface_data, interface_type - ) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract Interfaces data, Error: {ex}, Line: {line_number}" - ) - - def _extract_interface( - self, name: str, data: dict, interface_type: str | None = None - ) -> EdgeOSInterfaceData: - interface = self._interfaces.get(name) - - try: - if data is not None: - if interface is None: - interface = EdgeOSInterfaceData(name) - interface.set_type(interface_type) - - if interface.handler == InterfaceHandlers.IGNORED: - message = f"Interface {name} is ignored, no entities will be created, Data: {data}" - self.unique_log(logging.INFO, message) - - interface.description = data.get(INTERFACE_DATA_DESCRIPTION) - interface.duplex = data.get(INTERFACE_DATA_DUPLEX) - interface.speed = data.get(INTERFACE_DATA_SPEED) - interface.bridge_group = data.get(INTERFACE_DATA_BRIDGE_GROUP) - interface.address = data.get(INTERFACE_DATA_ADDRESS) - interface.aging = data.get(INTERFACE_DATA_AGING) - interface.bridged_conntrack = data.get(INTERFACE_DATA_BRIDGED_CONNTRACK) - interface.hello_time = data.get(INTERFACE_DATA_HELLO_TIME) - interface.max_age = data.get(INTERFACE_DATA_MAX_AGE) - interface.priority = data.get(INTERFACE_DATA_PRIORITY) - interface.promiscuous = data.get(INTERFACE_DATA_PROMISCUOUS) - interface.stp = ( - data.get(INTERFACE_DATA_STP, FALSE_STR).lower() == TRUE_STR - ) - - self._interfaces[interface.unique_id] = interface - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract interface data for {name}/{interface_type}, " - f"Error: {ex}, " - f"Line: {line_number}" - ) - - return interface - - @staticmethod - def _update_interface_stats(interface: EdgeOSInterfaceData, data: dict): - try: - if data is not None: - interface.up = ( - str(data.get(INTERFACE_DATA_UP, False)).lower() == TRUE_STR - ) - interface.l1up = ( - str(data.get(INTERFACE_DATA_LINK_UP, False)).lower() == TRUE_STR - ) - interface.mac = data.get(INTERFACE_DATA_MAC) - interface.multicast = data.get(INTERFACE_DATA_MULTICAST, 0) - interface.address = data.get(ADDRESS_LIST, []) - - directions = [interface.received, interface.sent] - - for direction in directions: - stat_data = {} - for stat_key in TRAFFIC_DATA_INTERFACE_ITEMS: - key = f"{direction.direction}_{stat_key}" - stat_data_item = TRAFFIC_DATA_INTERFACE_ITEMS.get(stat_key) - - stat_data[stat_data_item] = float(data.get(key)) - - direction.update(stat_data) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to update interface statistics for {interface.name}, " - f"Error: {ex}, " - f"Line: {line_number}" - ) - - @staticmethod - def _update_device_stats(device_data: EdgeOSDeviceData, data: dict): - try: - if not device_data.is_leased: - stats = [device_data.received, device_data.sent] - - for stat in stats: - stat_data = {} - for stat_key in TRAFFIC_DATA_DEVICE_ITEMS: - key = f"{stat.direction}_{stat_key}" - stat_data_item = TRAFFIC_DATA_DEVICE_ITEMS.get(stat_key) - - stat_data[stat_data_item] = data.get(key) - - stat.update(stat_data) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to update device statistics for {device_data.hostname}, " - f"Error: {ex}, " - f"Line: {line_number}" - ) - - def _update_system_stats(self, system_stats_data: dict, discovery_data: dict): - try: - system_data = self._system - - system_data.fw_version = discovery_data.get(DISCOVER_DATA_FW_VERSION) - system_data.product = discovery_data.get(DISCOVER_DATA_PRODUCT) - - uptime = float(system_stats_data.get(SYSTEM_STATS_DATA_UPTIME, 0)) - - system_data.cpu = int(system_stats_data.get(SYSTEM_STATS_DATA_CPU, 0)) - system_data.mem = int(system_stats_data.get(SYSTEM_STATS_DATA_MEM, 0)) - - if uptime != system_data.uptime: - system_data.uptime = uptime - system_data.last_reset = self._get_last_reset(uptime) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to update system statistics, " - f"Error: {ex}, " - f"Line: {line_number}" - ) - - def _extract_unknown_devices(self): - try: - unknown_devices = 0 - data_leases_stats = self.api.data.get(API_DATA_DHCP_STATS, {}) - - subnets = data_leases_stats.get(DHCP_SERVER_STATS, {}) - - for subnet in subnets: - subnet_data = subnets.get(subnet, {}) - unknown_devices += int(subnet_data.get(DHCP_SERVER_LEASED, 0)) - - self._system.leased_devices = unknown_devices - - data_leases = self.api.data.get(API_DATA_DHCP_LEASES, {}) - data_server_leases = data_leases.get(DHCP_SERVER_LEASES, {}) - - for subnet in data_server_leases: - subnet_data = data_server_leases.get(subnet, {}) - - for ip in subnet_data: - device_data = subnet_data.get(ip) - - hostname = device_data.get(DHCP_SERVER_LEASES_CLIENT_HOSTNAME) - - static_mapping_data = { - DHCP_SERVER_IP_ADDRESS: ip, - DHCP_SERVER_MAC_ADDRESS: device_data.get(DEVICE_DATA_MAC), - } - - self._set_device(hostname, None, static_mapping_data, True) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract Unknown Devices data, Error: {ex}, Line: {line_number}" - ) - - def _extract_devices(self, data: dict): - try: - service = data.get(DATA_SYSTEM_SERVICE, {}) - dhcp_server = service.get(DATA_SYSTEM_SERVICE_DHCP_SERVER, {}) - shared_network_names = dhcp_server.get(DHCP_SERVER_SHARED_NETWORK_NAME, {}) - - for shared_network_name in shared_network_names: - shared_network_name_data = shared_network_names.get( - shared_network_name, {} - ) - subnets = shared_network_name_data.get(DHCP_SERVER_SUBNET, {}) - - for subnet in subnets: - subnet_data = subnets.get(subnet, {}) - - domain_name = subnet_data.get(SYSTEM_DATA_DOMAIN_NAME) - static_mappings = subnet_data.get(DHCP_SERVER_STATIC_MAPPING, {}) - - for hostname in static_mappings: - static_mapping_data = static_mappings.get(hostname, {}) - - self._set_device( - hostname, domain_name, static_mapping_data, False - ) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract Devices data, Error: {ex}, Line: {line_number}" - ) - - def _set_device( - self, - hostname: str, - domain_name: str | None, - static_mapping_data: dict, - is_leased: bool, - ): - ip_address = static_mapping_data.get(DHCP_SERVER_IP_ADDRESS) - mac_address = static_mapping_data.get(DHCP_SERVER_MAC_ADDRESS) - - existing_device_data = self._devices.get(mac_address) - - if existing_device_data is None: - device_data = EdgeOSDeviceData( - hostname, ip_address, mac_address, domain_name, is_leased - ) - - else: - device_data = existing_device_data - - self._devices[device_data.unique_id] = device_data - self._devices_ip_mapping[device_data.ip] = device_data.unique_id - - def _get_device(self, unique_id: str) -> EdgeOSDeviceData | None: - device = self._devices.get(unique_id) - - return device - - def _get_device_by_ip(self, ip: str) -> EdgeOSDeviceData | None: - unique_id = self._devices_ip_mapping.get(ip) - - device = self._get_device(unique_id) - - return device - - def _set_ha_device( - self, name: str, model: str, manufacturer: str, version: str | None = None - ): - device_details = self.device_manager.get(name) - - device_details_data = { - "identifiers": {(DEFAULT_NAME, name)}, - "name": name, - "manufacturer": manufacturer, - "model": model, - } - - if version is not None: - device_details_data["sw_version"] = version - - if device_details is None or device_details != device_details_data: - self.device_manager.set(name, device_details_data) - - _LOGGER.debug(f"Created HA device {name} [{model}]") - - def _load_main_device(self): - self._set_ha_device( - self.system_name, - self._system.product, - MANUFACTURER, - self._system.fw_version, - ) - - def _load_device_device(self, device: EdgeOSDeviceData): - name = self.get_device_name(device) - self._set_ha_device(name, "Device", DEFAULT_NAME) - - def _load_interface_device(self, interface: EdgeOSInterfaceData): - name = self.get_interface_name(interface) - self._set_ha_device(name, "Interface", DEFAULT_NAME) - - def _load_unknown_devices_sensor(self): - device_name = self.system_name - entity_name = f"{device_name} Unknown Devices" - - try: - state = self._system.leased_devices - - leased_devices = [] - - for unique_id in self._devices: - device = self._devices.get(unique_id) - - if device.is_leased: - leased_devices.append(f"{device.hostname} ({device.ip})") - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - DHCP_SERVER_LEASED: leased_devices, - } - - unique_id = EntityData.generate_unique_id(DOMAIN_SENSOR, entity_name) - icon = "mdi:help-network-outline" - - entity_description = SensorEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - state_class=SensorStateClass.MEASUREMENT, - ) - - self.entity_manager.set_entity( - DOMAIN_SENSOR, - self.entry_id, - state, - attributes, - device_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load sensor for {entity_name}") - - def _load_cpu_sensor(self): - device_name = self.system_name - entity_name = f"{device_name} CPU Usage" - - try: - state = self._system.cpu - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - } - - unique_id = EntityData.generate_unique_id(DOMAIN_SENSOR, entity_name) - icon = "mdi:chip" - - entity_description = SensorEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - state_class=SensorStateClass.MEASUREMENT, - native_unit_of_measurement="%", - ) - - self.entity_manager.set_entity( - DOMAIN_SENSOR, - self.entry_id, - state, - attributes, - device_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load sensor for {entity_name}") - - def _load_ram_sensor(self): - device_name = self.system_name - entity_name = f"{device_name} RAM Usage" - - try: - state = self._system.mem - - attributes = {ATTR_FRIENDLY_NAME: entity_name} - - unique_id = EntityData.generate_unique_id(DOMAIN_SENSOR, entity_name) - icon = "mdi:memory" - - entity_description = SensorEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - state_class=SensorStateClass.MEASUREMENT, - native_unit_of_measurement="%", - ) - - self.entity_manager.set_entity( - DOMAIN_SENSOR, - self.entry_id, - state, - attributes, - device_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load sensor for {entity_name}") - - def _load_uptime_sensor(self): - device_name = self.system_name - entity_name = f"{device_name} Last Restart" - - try: - state = self._system.uptime - - attributes = {ATTR_FRIENDLY_NAME: entity_name} - - unique_id = EntityData.generate_unique_id(DOMAIN_SENSOR, entity_name) - icon = "mdi:credit-card-clock" - - entity_description = SensorEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - state_class=SensorStateClass.TOTAL_INCREASING, - ) - - self.entity_manager.set_entity( - DOMAIN_SENSOR, - self.entry_id, - state, - attributes, - device_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load sensor for {entity_name}") - - def _load_firmware_upgrade_binary_sensor(self): - device_name = self.system_name - entity_name = f"{device_name} Firmware Upgrade" - - try: - state = STATE_ON if self._system.upgrade_available else STATE_OFF - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - SYSTEM_INFO_DATA_FW_LATEST_URL: self._system.upgrade_url, - SYSTEM_INFO_DATA_FW_LATEST_VERSION: self._system.upgrade_version, - } - - unique_id = EntityData.generate_unique_id(DOMAIN_BINARY_SENSOR, entity_name) - - entity_description = BinarySensorEntityDescription( - key=unique_id, - name=entity_name, - device_class=BinarySensorDeviceClass.UPDATE, - ) - - self.entity_manager.set_entity( - DOMAIN_BINARY_SENSOR, - self.entry_id, - state, - attributes, - device_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load sensor for {entity_name}") - - def _load_log_incoming_messages_switch(self): - device_name = self.system_name - entity_name = f"{device_name} Log Incoming Messages" - - try: - state = self.storage_api.log_incoming_messages - - attributes = {ATTR_FRIENDLY_NAME: entity_name} - - unique_id = EntityData.generate_unique_id(DOMAIN_SWITCH, entity_name) - - icon = "mdi:math-log" - - entity_description = SwitchEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - entity_category=EntityCategory.CONFIG, - ) - - self.entity_manager.set_entity( - DOMAIN_SWITCH, - self.entry_id, - state, - attributes, - device_name, - entity_description, - ) - - self.set_action( - unique_id, - ACTION_CORE_ENTITY_TURN_ON, - self._enable_log_incoming_messages, - ) - self.set_action( - unique_id, - ACTION_CORE_ENTITY_TURN_OFF, - self._disable_log_incoming_messages, - ) - - except Exception as ex: - self.log_exception( - ex, f"Failed to load log incoming messages switch for {entity_name}" - ) - - def _load_device_tracker(self, device: EdgeOSDeviceData): - device_name = self.get_device_name(device) - entity_name = f"{device_name}" - - try: - state = ( - device.last_activity_in_seconds - <= self.storage_api.consider_away_interval - ) - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - } - - unique_id = EntityData.generate_unique_id( - DOMAIN_DEVICE_TRACKER, entity_name - ) - - entity_description = EntityDescription(key=unique_id, name=entity_name) - - details = {ENTITY_UNIQUE_ID: device.unique_id} - - is_monitored = self.storage_api.monitored_devices.get( - device.unique_id, False - ) - - self.entity_manager.set_entity( - DOMAIN_DEVICE_TRACKER, - self.entry_id, - state, - attributes, - device_name, - entity_description, - destructors=[not is_monitored], - details=details, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load device tracker for {entity_name}") - - def _load_device_monitor_switch(self, device: EdgeOSDeviceData): - device_name = self.get_device_name(device) - entity_name = f"{device_name} Monitored" - - try: - state = self.storage_api.monitored_devices.get(device.unique_id, False) - - attributes = {ATTR_FRIENDLY_NAME: entity_name} - - unique_id = EntityData.generate_unique_id(DOMAIN_SWITCH, entity_name) - icon = "mdi:monitor-eye" - - entity_description = SwitchEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - entity_category=EntityCategory.CONFIG, - ) - - details = {ENTITY_UNIQUE_ID: device.unique_id} - - self.set_action( - unique_id, ACTION_CORE_ENTITY_TURN_ON, self._set_device_monitored - ) - self.set_action( - unique_id, ACTION_CORE_ENTITY_TURN_OFF, self._set_device_unmonitored - ) - - self.entity_manager.set_entity( - DOMAIN_SWITCH, - self.entry_id, - state, - attributes, - device_name, - entity_description, - details=details, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load switch for {entity_name}") - - def _load_stats_sensor( - self, - item_unique_id: str, - device_name: str, - entity_suffix: str, - state: str | int | float | None, - monitored_items: dict, - ): - entity_name = f"{device_name} {entity_suffix}" - - try: - attributes = {ATTR_FRIENDLY_NAME: entity_name} - unique_id = EntityData.generate_unique_id(DOMAIN_SENSOR, entity_name) - - icon = STATS_ICONS.get(entity_suffix) - is_monitored = monitored_items.get(item_unique_id, False) - - device_class: SensorDeviceClass | None = None - state_class = SensorStateClass.MEASUREMENT - - if entity_suffix in STATS_DATA_RATE: - unit_of_measurement = UnitOfDataRate.BYTES_PER_SECOND - device_class = SensorDeviceClass.DATA_RATE - - elif entity_suffix in STATS_DATA_SIZE: - unit_of_measurement = UnitOfInformation.BYTES - device_class = SensorDeviceClass.DATA_SIZE - state_class = SensorStateClass.TOTAL_INCREASING - - else: - unit_of_measurement = str(STATS_UNITS.get(entity_suffix)).capitalize() - - entity_description = SensorEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - state_class=state_class, - native_unit_of_measurement=unit_of_measurement, - device_class=device_class, - ) - - if unit_of_measurement.lower() in [ - TRAFFIC_DATA_ERRORS, - TRAFFIC_DATA_PACKETS, - TRAFFIC_DATA_DROPPED, - ]: - state = self._format_number(state) - - self.entity_manager.set_entity( - DOMAIN_SENSOR, - self.entry_id, - state, - attributes, - device_name, - entity_description, - destructors=[not is_monitored], - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load sensor for {entity_name}") - - def _load_interface_status_switch(self, interface: EdgeOSInterfaceData): - interface_name = self.get_interface_name(interface) - entity_name = f"{interface_name} Status" - - try: - state = interface.up - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - ADDRESS_LIST: interface.address, - } - - unique_id = EntityData.generate_unique_id(DOMAIN_SWITCH, entity_name) - icon = "mdi:eye-settings" - - entity_description = SwitchEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - entity_category=EntityCategory.CONFIG, - ) - - details = {ENTITY_UNIQUE_ID: interface.unique_id} - - self.set_action( - unique_id, ACTION_CORE_ENTITY_TURN_ON, self._set_interface_enabled - ) - self.set_action( - unique_id, ACTION_CORE_ENTITY_TURN_OFF, self._set_interface_disabled - ) - - self.entity_manager.set_entity( - DOMAIN_SWITCH, - self.entry_id, - state, - attributes, - interface_name, - entity_description, - details=details, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load switch for {entity_name}") - - def _load_interface_status_binary_sensor(self, interface: EdgeOSInterfaceData): - interface_name = self.get_interface_name(interface) - entity_name = f"{interface_name} Status" - - try: - state = STATE_ON if interface.up else STATE_OFF - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - ADDRESS_LIST: interface.address, - } - - unique_id = EntityData.generate_unique_id(DOMAIN_BINARY_SENSOR, entity_name) - - entity_description = BinarySensorEntityDescription( - key=unique_id, - name=entity_name, - device_class=BinarySensorDeviceClass.CONNECTIVITY, - ) - - self.entity_manager.set_entity( - DOMAIN_BINARY_SENSOR, - self.entry_id, - state, - attributes, - interface_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load binary sensor for {entity_name}") - - def _load_interface_connected_binary_sensor(self, interface: EdgeOSInterfaceData): - interface_name = self.get_interface_name(interface) - entity_name = f"{interface_name} Connected" - - try: - state = STATE_ON if interface.l1up else STATE_OFF - - attributes = { - ATTR_FRIENDLY_NAME: entity_name, - ADDRESS_LIST: interface.address, - } - - unique_id = EntityData.generate_unique_id(DOMAIN_BINARY_SENSOR, entity_name) - - entity_description = BinarySensorEntityDescription( - key=unique_id, - name=entity_name, - device_class=BinarySensorDeviceClass.CONNECTIVITY, - ) - - self.entity_manager.set_entity( - DOMAIN_BINARY_SENSOR, - self.entry_id, - state, - attributes, - interface_name, - entity_description, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load binary sensor for {entity_name}") - - def _load_interface_monitor_switch(self, interface: EdgeOSInterfaceData): - interface_name = self.get_interface_name(interface) - entity_name = f"{interface_name} Monitored" - - try: - state = self.storage_api.monitored_interfaces.get( - interface.unique_id, False - ) - - attributes = {ATTR_FRIENDLY_NAME: entity_name} - - unique_id = EntityData.generate_unique_id(DOMAIN_SWITCH, entity_name) - icon = None - - entity_description = SwitchEntityDescription( - key=unique_id, - name=entity_name, - icon=icon, - entity_category=EntityCategory.CONFIG, - ) - - details = {ENTITY_UNIQUE_ID: interface.unique_id} - - self.set_action( - unique_id, ACTION_CORE_ENTITY_TURN_ON, self._set_interface_monitored - ) - self.set_action( - unique_id, ACTION_CORE_ENTITY_TURN_OFF, self._set_interface_unmonitored - ) - - self.entity_manager.set_entity( - DOMAIN_SWITCH, - self.entry_id, - state, - attributes, - interface_name, - entity_description, - details=details, - ) - - except Exception as ex: - self.log_exception(ex, f"Failed to load switch for {entity_name}") - - async def _set_interface_enabled(self, entity: EntityData): - interface_item = self._get_interface_from_entity(entity) - - await self.api.set_interface_state(interface_item, True) - - async def _set_interface_disabled(self, entity: EntityData): - interface_item = self._get_interface_from_entity(entity) - - await self.api.set_interface_state(interface_item, False) - - async def _set_interface_monitored(self, entity: EntityData): - interface_item = self._get_interface_from_entity(entity) - - await self.storage_api.set_monitored_interface(interface_item.unique_id, True) - - async def _set_interface_unmonitored(self, entity: EntityData): - interface_item = self._get_interface_from_entity(entity) - - await self.storage_api.set_monitored_interface(interface_item.unique_id, False) - - async def _set_device_monitored(self, entity: EntityData): - device_item = self._get_device_from_entity(entity) - - await self.storage_api.set_monitored_device(device_item.unique_id, True) - - async def _set_device_unmonitored(self, entity: EntityData): - device_item = self._get_device_from_entity(entity) - - await self.storage_api.set_monitored_device(device_item.unique_id, False) - - async def _enable_log_incoming_messages(self, entity: EntityData): - await self.storage_api.set_log_incoming_messages(True) - - async def _disable_log_incoming_messages(self, entity: EntityData): - await self.storage_api.set_log_incoming_messages(False) - - def _get_device_from_entity(self, entity: EntityData) -> EdgeOSDeviceData: - unique_id = entity.details.get(ENTITY_UNIQUE_ID) - device_item = self._get_device(unique_id) - - return device_item - - def _get_interface_from_entity(self, entity: EntityData) -> EdgeOSInterfaceData: - unique_id = entity.details.get(ENTITY_UNIQUE_ID) - interface_item = self._interfaces.get(unique_id) - - return interface_item - - @staticmethod - def _format_number(value: int | float | None, digits: int = 0) -> int | float: - if value is None: - value = 0 - - value_str = f"{value:.{digits}f}" - result = int(value_str) if digits == 0 else float(value_str) - - return result - - async def _reload_integration(self): - data = {ENTITY_CONFIG_ENTRY_ID: self.entry_id} - - await self._hass.services.async_call(HA_NAME, SERVICE_RELOAD_CONFIG_ENTRY, data) - - def _update_configuration(self, service_call): - self._hass.async_create_task(self._async_update_configuration(service_call)) - - async def _async_update_configuration(self, service_call): - service_data = service_call.data - device_id = service_data.get(CONF_DEVICE_ID) - - _LOGGER.info(f"Update configuration called with data: {service_data}") - - if device_id is None: - _LOGGER.error("Operation cannot be performed, missing device information") - - else: - dr = async_get_device_registry(self._hass) - device = dr.devices.get(device_id) - can_handle_device = self.entry_id in device.config_entries - - if can_handle_device: - updated = await self._update_configuration_data(service_data) - - if updated: - await self._reload_integration() - - async def _update_configuration_data(self, data: dict): - result = False - - storage_data_import_keys: dict[ - str, Callable[[int | bool | str], Awaitable[None]] - ] = { - STORAGE_DATA_CONSIDER_AWAY_INTERVAL: self.storage_api.set_consider_away_interval, - STORAGE_DATA_UPDATE_ENTITIES_INTERVAL: self.storage_api.set_update_entities_interval, - STORAGE_DATA_UPDATE_API_INTERVAL: self.storage_api.set_update_api_interval, - STORAGE_DATA_LOG_INCOMING_MESSAGES: self.storage_api.set_log_incoming_messages, - } - - for key in storage_data_import_keys: - data_item = data.get(key.replace(STRING_DASH, STRING_UNDERSCORE)) - existing_data = self.storage_api.data.get(key) - - if data_item is not None and data_item != existing_data: - if not result: - result = True - - set_func = storage_data_import_keys.get(key) - - await set_func(data_item) - - return result - - @staticmethod - def _get_last_reset(uptime): - now = datetime.now().timestamp() - last_reset = int(now) - uptime - - result = datetime.fromtimestamp(last_reset) - - return result - - def unique_log(self, log_level: int, message: str): - if message not in self._unique_messages: - self._unique_messages.append(message) - - _LOGGER.log(log_level, self._unique_messages) diff --git a/custom_components/edgeos/component/models/__init__.py b/custom_components/edgeos/component/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/component/models/exceptions.py b/custom_components/edgeos/component/models/exceptions.py deleted file mode 100644 index 7dadebf..0000000 --- a/custom_components/edgeos/component/models/exceptions.py +++ /dev/null @@ -1,22 +0,0 @@ -from homeassistant.exceptions import HomeAssistantError - - -class IncompatibleVersion(HomeAssistantError): - def __init__(self, version): - self._version = version - - def __repr__(self): - return f"Unsupported EdgeOS version ({self._version})" - - -class SessionTerminatedException(HomeAssistantError): - Terminated = True - - -class LoginException(HomeAssistantError): - def __init__(self, status_code): - self._status_code = status_code - - @property - def status_code(self): - return self._status_code diff --git a/custom_components/edgeos/config_flow.py b/custom_components/edgeos/config_flow.py index 50fabcd..07d7c21 100644 --- a/custom_components/edgeos/config_flow.py +++ b/custom_components/edgeos/config_flow.py @@ -3,17 +3,12 @@ import logging -from cryptography.fernet import InvalidToken -import voluptuous as vol - from homeassistant import config_entries from homeassistant.config_entries import ConfigEntry from homeassistant.core import callback -from .component.api.api import IntegrationAPI -from .configuration.helpers.const import DEFAULT_NAME, DOMAIN -from .configuration.helpers.exceptions import AlreadyExistsError, LoginError -from .configuration.managers.configuration_manager import ConfigurationManager +from .common.consts import DOMAIN +from .managers.flow_manager import IntegrationFlowManager _LOGGER = logging.getLogger(__name__) @@ -28,8 +23,6 @@ class DomainFlowHandler(config_entries.ConfigFlow): def __init__(self): super().__init__() - self._config_manager: ConfigurationManager | None = None - @staticmethod @callback def async_get_options_flow(config_entry): @@ -38,94 +31,24 @@ def async_get_options_flow(config_entry): async def async_step_user(self, user_input=None): """Handle a flow start.""" - _LOGGER.debug(f"Starting async_step_user of {DEFAULT_NAME}") - - api = IntegrationAPI(self.hass) - self._config_manager = ConfigurationManager(self.hass, api) - - await self._config_manager.initialize() - - errors = None - - if user_input is not None: - try: - await self._config_manager.validate(user_input) - - _LOGGER.debug("User inputs are valid") - - return self.async_create_entry(title=DEFAULT_NAME, data=user_input) - except LoginError as lex: - errors = lex.errors - - except InvalidToken: - errors = {"base": "corrupted_encryption_key"} - - except AlreadyExistsError: - errors = {"base": "already_configured"} - - if errors is not None: - error_message = errors.get("base") - - _LOGGER.warning(f"Failed to create integration, Error: {error_message}") - - new_user_input = self._config_manager.get_data_fields(user_input) - - schema = vol.Schema(new_user_input) + flow_manager = IntegrationFlowManager(self.hass, self) - return self.async_show_form(step_id="user", data_schema=schema, errors=errors) + return await flow_manager.async_step(user_input) class DomainOptionsFlowHandler(config_entries.OptionsFlow): """Handle domain options.""" + _config_entry: ConfigEntry + def __init__(self, config_entry: ConfigEntry): """Initialize domain options flow.""" super().__init__() self._config_entry = config_entry - self._config_manager: ConfigurationManager | None = None async def async_step_init(self, user_input=None): """Manage the domain options.""" - _LOGGER.info(f"Starting additional settings step: {user_input}") - - api = IntegrationAPI(self.hass) - self._config_manager = ConfigurationManager(self.hass, api) - await self._config_manager.initialize() - - errors = None - - if user_input is not None: - try: - await self._config_manager.validate(user_input) - - _LOGGER.debug("User inputs are valid") - - options = self._config_manager.remap_entry_data( - self._config_entry, user_input - ) - - return self.async_create_entry( - title=self._config_entry.title, data=options - ) - except LoginError as lex: - errors = lex.errors - - except InvalidToken: - errors = {"base": "corrupted_encryption_key"} - - except AlreadyExistsError: - errors = {"base": "already_configured"} - - if errors is not None: - error_message = errors.get("base") - - _LOGGER.warning(f"Failed to create integration, Error: {error_message}") - - new_user_input = self._config_manager.get_options_fields( - self._config_entry.data - ) - - schema = vol.Schema(new_user_input) + flow_manager = IntegrationFlowManager(self.hass, self, self._config_entry) - return self.async_show_form(step_id="init", data_schema=schema, errors=errors) + return await flow_manager.async_step(user_input) diff --git a/custom_components/edgeos/configuration/__init__.py b/custom_components/edgeos/configuration/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/configuration/helpers/__init__.py b/custom_components/edgeos/configuration/helpers/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/configuration/helpers/const.py b/custom_components/edgeos/configuration/helpers/const.py deleted file mode 100644 index 5055adf..0000000 --- a/custom_components/edgeos/configuration/helpers/const.py +++ /dev/null @@ -1,25 +0,0 @@ -""" -Following constants are mandatory for CORE: - DEFAULT_NAME - Full name for the title of the integration - DOMAIN - name of component, will be used as component's domain - SUPPORTED_PLATFORMS - list of supported HA components to initialize -""" - -from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME - -DOMAIN = "edgeos" -DEFAULT_NAME = "EdgeOS" -MANUFACTURER = "Ubiquiti" - -DATA_KEYS = [CONF_HOST, CONF_USERNAME, CONF_PASSWORD] - -MAXIMUM_RECONNECT = 3 - -API_URL_TEMPLATE = "https://{}" -WEBSOCKET_URL_TEMPLATE = "wss://{}/ws/stats" - -COOKIE_PHPSESSID = "PHPSESSID" -COOKIE_BEAKER_SESSION_ID = "beaker.session.id" -COOKIE_CSRF_TOKEN = "X-CSRF-TOKEN" - -HEADER_CSRF_TOKEN = "X-Csrf-token" diff --git a/custom_components/edgeos/configuration/helpers/exceptions.py b/custom_components/edgeos/configuration/helpers/exceptions.py deleted file mode 100644 index c06f7d0..0000000 --- a/custom_components/edgeos/configuration/helpers/exceptions.py +++ /dev/null @@ -1,15 +0,0 @@ -from homeassistant.exceptions import HomeAssistantError - - -class AlreadyExistsError(HomeAssistantError): - title: str - - def __init__(self, title: str): - self.title = title - - -class LoginError(HomeAssistantError): - errors: dict - - def __init__(self, errors): - self.errors = errors diff --git a/custom_components/edgeos/configuration/managers/configuration_manager.py b/custom_components/edgeos/configuration/managers/configuration_manager.py deleted file mode 100644 index 71d21aa..0000000 --- a/custom_components/edgeos/configuration/managers/configuration_manager.py +++ /dev/null @@ -1,148 +0,0 @@ -from __future__ import annotations - -import logging -import sys -from typing import Any - -import voluptuous as vol - -from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME -from homeassistant.core import HomeAssistant - -from ...core.api.base_api import BaseAPI -from ...core.helpers.enums import ConnectivityStatus -from ...core.managers.password_manager import PasswordManager -from ..helpers.const import DATA_KEYS -from ..helpers.exceptions import LoginError -from ..models.config_data import ConfigData - -_LOGGER = logging.getLogger(__name__) - - -class ConfigurationManager: - password_manager: PasswordManager - config: dict[str, ConfigData] - api: BaseAPI | None - - def __init__(self, hass: HomeAssistant, api: BaseAPI | None = None): - self.hass = hass - self.config = {} - self.password_manager = PasswordManager(hass) - self.api = api - - async def initialize(self): - await self.password_manager.initialize() - - def get(self, entry_id: str): - config = self.config.get(entry_id) - - return config - - async def load(self, entry: ConfigEntry): - try: - await self.initialize() - - config = {k: entry.data[k] for k in entry.data} - - if CONF_PASSWORD in config: - encrypted_password = config[CONF_PASSWORD] - - config[CONF_PASSWORD] = self.password_manager.get(encrypted_password) - - config_data = ConfigData.from_dict(config) - - if config_data is not None: - config_data.entry = entry - - self.config[entry.entry_id] = config_data - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to load configuration, error: {str(ex)}, line: {line_number}" - ) - - async def validate(self, data: dict[str, Any]): - if self.api is None: - _LOGGER.error("Validate configuration is not supported through that flow") - return - - _LOGGER.debug("Validate login") - - await self.api.validate(data) - - errors = self._get_config_errors() - - if errors is None: - password = data[CONF_PASSWORD] - - data[CONF_PASSWORD] = self.password_manager.set(password) - - else: - raise LoginError(errors) - - def _get_config_errors(self): - result = None - status_mapping = { - str(ConnectivityStatus.NotConnected): "invalid_server_details", - str(ConnectivityStatus.Connecting): "invalid_server_details", - str(ConnectivityStatus.Failed): "invalid_server_details", - str(ConnectivityStatus.NotFound): "invalid_server_details", - str(ConnectivityStatus.MissingAPIKey): "missing_permanent_api_key", - str(ConnectivityStatus.InvalidCredentials): "invalid_admin_credentials", - str(ConnectivityStatus.TemporaryConnected): "missing_permanent_api_key", - } - - status_description = status_mapping.get(str(self.api.status)) - - if status_description is not None: - result = {"base": status_description} - - return result - - @staticmethod - def get_data_fields(user_input: dict[str, Any] | None) -> dict[vol.Marker, Any]: - if user_input is None: - user_input = ConfigData.from_dict().to_dict() - - fields = { - vol.Required(CONF_HOST, default=user_input.get(CONF_HOST)): str, - vol.Required(CONF_USERNAME, default=user_input.get(CONF_USERNAME)): str, - vol.Required(CONF_PASSWORD, default=user_input.get(CONF_PASSWORD)): str, - } - - return fields - - def get_options_fields(self, user_input: dict[str, Any]) -> dict[vol.Marker, Any]: - if user_input is None: - data = ConfigData.from_dict().to_dict() - - else: - data = {k: user_input[k] for k in user_input} - encrypted_password = data.get(CONF_PASSWORD) - - data[CONF_PASSWORD] = self.password_manager.get(encrypted_password) - - fields = self.get_data_fields(data) - - return fields - - def remap_entry_data( - self, entry: ConfigEntry, options: dict[str, Any] - ) -> dict[str, Any]: - config_options = {} - config_data = {} - - for key in options: - if key in DATA_KEYS: - config_data[key] = options.get(key, entry.data.get(key)) - - else: - config_options[key] = options.get(key) - - config_entries = self.hass.config_entries - config_entries.async_update_entry(entry, data=config_data) - - return config_options diff --git a/custom_components/edgeos/configuration/models/__init__.py b/custom_components/edgeos/configuration/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/configuration/models/config_data.py b/custom_components/edgeos/configuration/models/config_data.py deleted file mode 100644 index 5ea9ad2..0000000 --- a/custom_components/edgeos/configuration/models/config_data.py +++ /dev/null @@ -1,52 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME - -from ..helpers.const import API_URL_TEMPLATE - - -class ConfigData: - host: str | None - username: str | None - password: str | None - entry: ConfigEntry | None - - def __init__(self): - self.host = None - self.username = None - self.password = None - self.entry = None - - @property - def url(self): - url = API_URL_TEMPLATE.format(self.host) - - return url - - @staticmethod - def from_dict(data: dict[str, Any] = None) -> ConfigData: - result = ConfigData() - - if data is not None: - result.host = data.get(CONF_HOST) - result.username = data.get(CONF_USERNAME) - result.password = data.get(CONF_PASSWORD) - - return result - - def to_dict(self): - obj = { - CONF_HOST: self.host, - CONF_USERNAME: self.username, - CONF_PASSWORD: self.password, - } - - return obj - - def __repr__(self): - to_string = f"{self.to_dict()}" - - return to_string diff --git a/custom_components/edgeos/core/__init__.py b/custom_components/edgeos/core/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/core/api/__init__.py b/custom_components/edgeos/core/api/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/core/api/base_api.py b/custom_components/edgeos/core/api/base_api.py deleted file mode 100644 index 53c6035..0000000 --- a/custom_components/edgeos/core/api/base_api.py +++ /dev/null @@ -1,93 +0,0 @@ -from __future__ import annotations - -from collections.abc import Awaitable, Callable -import logging -import sys - -from aiohttp import ClientSession - -from homeassistant.core import HomeAssistant -from homeassistant.helpers.aiohttp_client import async_create_clientsession - -from ..helpers.enums import ConnectivityStatus - -_LOGGER = logging.getLogger(__name__) - - -class BaseAPI: - """The Class for handling the data retrieval.""" - - hass: HomeAssistant - session: ClientSession | None - status: ConnectivityStatus - data: dict - onDataChangedAsync: Callable[[], Awaitable[None]] | None = None - onStatusChangedAsync: Callable[[ConnectivityStatus], Awaitable[None]] | None = None - - def __init__( - self, - hass: HomeAssistant | None, - async_on_data_changed: Callable[[], Awaitable[None]] | None = None, - async_on_status_changed: Callable[[ConnectivityStatus], Awaitable[None]] - | None = None, - ): - self.hass = hass - self.status = ConnectivityStatus.NotConnected - self.data = {} - self.onDataChangedAsync = async_on_data_changed - self.onStatusChangedAsync = async_on_status_changed - - self.session = None - - @property - def is_home_assistant(self): - return self.hass is not None - - async def initialize_session(self, cookies=None, cookie_jar=None): - try: - if self.is_home_assistant: - self.session = async_create_clientsession( - hass=self.hass, cookies=cookies, cookie_jar=cookie_jar - ) - - else: - self.session = ClientSession(cookies=cookies, cookie_jar=cookie_jar) - - await self.login() - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.warning( - f"Failed to initialize session, Error: {str(ex)}, Line: {line_number}" - ) - - await self.set_status(ConnectivityStatus.Failed) - - async def login(self): - _LOGGER.info("Performing login") - - await self.set_status(ConnectivityStatus.Connecting) - - async def validate(self, data: dict | None = None): - pass - - async def terminate(self): - self.data = {} - - await self.set_status(ConnectivityStatus.Disconnected) - - async def set_status(self, status: ConnectivityStatus): - if status != self.status: - self.status = status - - await self.fire_status_changed_event() - - async def fire_status_changed_event(self): - if self.onStatusChangedAsync is not None: - await self.onStatusChangedAsync(self.status) - - async def fire_data_changed_event(self): - if self.onDataChangedAsync is not None: - await self.onDataChangedAsync() diff --git a/custom_components/edgeos/core/components/__init__.py b/custom_components/edgeos/core/components/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/core/components/binary_sensor.py b/custom_components/edgeos/core/components/binary_sensor.py deleted file mode 100644 index 979292e..0000000 --- a/custom_components/edgeos/core/components/binary_sensor.py +++ /dev/null @@ -1,32 +0,0 @@ -""" -Support for binary sensors. -""" -from __future__ import annotations - -from homeassistant.components.binary_sensor import BinarySensorEntity -from homeassistant.const import STATE_ON -from homeassistant.core import HomeAssistant - -from ..helpers.const import DOMAIN_BINARY_SENSOR -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - - -class CoreBinarySensor(BinarySensorEntity, BaseEntity): - """Representation a binary sensor that is updated.""" - - @property - def is_on(self): - """Return true if the binary sensor is on.""" - return self.entity.state == STATE_ON - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - binary_sensor = CoreBinarySensor() - binary_sensor.initialize(hass, entity, DOMAIN_BINARY_SENSOR) - - return binary_sensor - - @staticmethod - def get_domain(): - return DOMAIN_BINARY_SENSOR diff --git a/custom_components/edgeos/core/components/camera.py b/custom_components/edgeos/core/components/camera.py deleted file mode 100644 index 7ecba4f..0000000 --- a/custom_components/edgeos/core/components/camera.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -Support for camera. -""" -from __future__ import annotations - -from abc import ABC -import asyncio -from datetime import datetime -import logging -import sys - -import aiohttp -import async_timeout - -from homeassistant.components.camera import ( - DEFAULT_CONTENT_TYPE, - Camera, - CameraEntityFeature, -) -from homeassistant.core import HomeAssistant -from homeassistant.exceptions import TemplateError -from homeassistant.helpers.aiohttp_client import async_get_clientsession -import homeassistant.helpers.config_validation as cv - -from ..helpers.const import ( - ATTR_MODE_RECORD, - ATTR_STREAM_FPS, - CONF_MOTION_DETECTION, - CONF_STILL_IMAGE_URL, - CONF_STREAM_SOURCE, - DOMAIN_CAMERA, - EMPTY_STRING, - SINGLE_FRAME_PS, -) -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class CoreCamera(Camera, BaseEntity, ABC): - """Camera""" - - def __init__(self, hass, device_info): - super().__init__() - self.hass = hass - self._still_image_url = None - self._stream_source = None - self._frame_interval = 0 - self._supported_features = 0 - self.content_type = DEFAULT_CONTENT_TYPE - self._auth = None - self._last_url = None - self._last_image = None - self.verify_ssl = False - self._is_recording_state = None - - def initialize( - self, - hass: HomeAssistant, - entity: EntityData, - current_domain: str, - DOMAIN_STREAM=None, - ): - super().initialize(hass, entity, current_domain) - - try: - if self.ha is None: - _LOGGER.warning("Failed to initialize CoreCamera without HA manager") - return - - config_data = self.ha.config_data - - username = config_data.username - password = config_data.password - - fps_str = entity.details.get(ATTR_STREAM_FPS, SINGLE_FRAME_PS) - - fps = SINGLE_FRAME_PS if fps_str == EMPTY_STRING else int(float(fps_str)) - - stream_source = entity.attributes.get(CONF_STREAM_SOURCE) - - snapshot = entity.attributes.get(CONF_STILL_IMAGE_URL) - - still_image_url_template = cv.template(snapshot) - - self._still_image_url = still_image_url_template - self._still_image_url.hass = hass - - self._stream_source = stream_source - self._frame_interval = SINGLE_FRAME_PS / fps - - self._is_recording_state = self.entity.details.get(ATTR_MODE_RECORD) - self._attr_is_streaming = stream_source is not None - - if self._stream_source: - self._attr_supported_features = CameraEntityFeature.STREAM - else: - self._attr_supported_features = CameraEntityFeature(0) - - if username and password: - self._auth = aiohttp.BasicAuth(username, password=password) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to initialize CoreCamera instance, Error: {ex}, Line: {line_number}" - ) - - @property - def is_recording(self) -> bool: - return self.entity.state == self._is_recording_state - - @property - def motion_detection_enabled(self): - return self.entity.details.get(CONF_MOTION_DETECTION, False) - - @property - def frame_interval(self): - """Return the interval between frames of the mjpeg stream.""" - return self._frame_interval - - def camera_image( - self, width: int | None = None, height: int | None = None - ) -> bytes | None: - """Return bytes of camera image.""" - return asyncio.run_coroutine_threadsafe( - self.async_camera_image(), self.hass.loop - ).result() - - async def async_camera_image( - self, width: int | None = None, height: int | None = None - ) -> bytes | None: - """Return a still image response from the camera.""" - try: - url = self._still_image_url.async_render() - except TemplateError as err: - _LOGGER.error( - f"Error parsing template {self._still_image_url}, Error: {err}" - ) - return self._last_image - - try: - ws = async_get_clientsession(self.hass, verify_ssl=self.verify_ssl) - async with async_timeout.timeout(10): - url = f"{url}?ts={datetime.now().timestamp()}" - response = await ws.get(url, auth=self._auth) - - self._last_image = await response.read() - - except asyncio.TimeoutError: - _LOGGER.error(f"Timeout getting camera image from {self.name}") - return self._last_image - - except aiohttp.ClientError as err: - _LOGGER.error( - f"Error getting new camera image from {self.name}, Error: {err}" - ) - return self._last_image - - self._last_url = url - return self._last_image - - async def stream_source(self): - """Return the source of the stream.""" - return self._stream_source - - async def async_enable_motion_detection(self) -> None: - """Enable motion detection in the camera.""" - if self.motion_detection_enabled: - _LOGGER.error(f"{self.name} - motion detection already enabled'") - - else: - await self.ha.async_core_entity_enable_motion_detection(self.entity) - - async def async_disable_motion_detection(self) -> None: - """Disable motion detection in camera.""" - if self.motion_detection_enabled: - await self.ha.async_core_entity_disable_motion_detection(self.entity) - - else: - _LOGGER.error(f"{self.name} - motion detection already disabled'") - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - camera = CoreCamera(hass, entity.details) - camera.initialize(hass, entity, DOMAIN_CAMERA) - - return camera - - @staticmethod - def get_domain(): - return DOMAIN_CAMERA diff --git a/custom_components/edgeos/core/components/device_tracker.py b/custom_components/edgeos/core/components/device_tracker.py deleted file mode 100644 index bb20907..0000000 --- a/custom_components/edgeos/core/components/device_tracker.py +++ /dev/null @@ -1,52 +0,0 @@ -""" -Support for device tracker. -""" -from __future__ import annotations - -import logging - -from homeassistant.components.device_tracker import ATTR_SOURCE_TYPE -from homeassistant.components.device_tracker.config_entry import ScannerEntity -from homeassistant.components.device_tracker.const import ATTR_IP, ATTR_MAC, SourceType -from homeassistant.core import HomeAssistant - -from ..helpers.const import DOMAIN_DEVICE_TRACKER -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class CoreScanner(BaseEntity, ScannerEntity): - """Represent a tracked device.""" - - @property - def ip_address(self) -> str | None: - """Return the primary ip address of the device.""" - return self.entity.details.get(ATTR_IP) - - @property - def mac_address(self) -> str | None: - """Return the mac address of the device.""" - return self.entity.details.get(ATTR_MAC) - - @property - def is_connected(self): - """Return true if the device is connected to the network.""" - return self.entity.state - - @property - def source_type(self): - """Return the source type.""" - return self.entity.attributes.get(ATTR_SOURCE_TYPE, SourceType.ROUTER) - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - device_tracker = CoreScanner() - device_tracker.initialize(hass, entity, DOMAIN_DEVICE_TRACKER) - - return device_tracker - - @staticmethod - def get_domain(): - return DOMAIN_DEVICE_TRACKER diff --git a/custom_components/edgeos/core/components/light.py b/custom_components/edgeos/core/components/light.py deleted file mode 100644 index 769269c..0000000 --- a/custom_components/edgeos/core/components/light.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -Support for light. -""" -from __future__ import annotations - -from abc import ABC -import logging -from typing import Any - -from homeassistant.components.light import ColorMode, LightEntity -from homeassistant.core import HomeAssistant - -from ..helpers.const import DOMAIN_LIGHT -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class CoreLight(LightEntity, BaseEntity, ABC): - """Class for a light.""" - - @property - def is_on(self) -> bool | None: - """Return the boolean response if the node is on.""" - return self.entity.state - - @property - def supported_color_modes(self) -> set[ColorMode] | set[str] | None: - """Flag supported color modes.""" - return set(ColorMode.ONOFF) - - async def async_turn_on(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_turn_on(self.entity) - - async def async_turn_off(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_turn_off(self.entity) - - async def async_setup(self): - pass - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - switch = CoreLight() - switch.initialize(hass, entity, DOMAIN_LIGHT) - - return switch - - @staticmethod - def get_domain(): - return DOMAIN_LIGHT diff --git a/custom_components/edgeos/core/components/select.py b/custom_components/edgeos/core/components/select.py deleted file mode 100644 index 38f42d4..0000000 --- a/custom_components/edgeos/core/components/select.py +++ /dev/null @@ -1,61 +0,0 @@ -""" -Support for select. -""" -from __future__ import annotations - -from abc import ABC -import logging -import sys - -from homeassistant.components.select import SelectEntity -from homeassistant.core import HomeAssistant - -from ..helpers.const import ATTR_OPTIONS, DOMAIN_SELECT -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class CoreSelect(SelectEntity, BaseEntity, ABC): - """Core Select""" - - def initialize( - self, - hass: HomeAssistant, - entity: EntityData, - current_domain: str, - ): - super().initialize(hass, entity, current_domain) - - try: - if hasattr(self.entity_description, ATTR_OPTIONS): - self._attr_options = getattr(self.entity_description, ATTR_OPTIONS) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to initialize CoreSelect instance, Error: {ex}, Line: {line_number}" - ) - - @property - def current_option(self) -> str: - """Return current lamp mode.""" - return str(self.entity.state) - - async def async_select_option(self, option: str) -> None: - """Select option.""" - await self.ha.async_core_entity_select_option(self.entity, option) - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - select = CoreSelect() - select.initialize(hass, entity, DOMAIN_SELECT) - - return select - - @staticmethod - def get_domain(): - return DOMAIN_SELECT diff --git a/custom_components/edgeos/core/components/sensor.py b/custom_components/edgeos/core/components/sensor.py deleted file mode 100644 index cc7dde1..0000000 --- a/custom_components/edgeos/core/components/sensor.py +++ /dev/null @@ -1,26 +0,0 @@ -from homeassistant.components.sensor import SensorEntity -from homeassistant.core import HomeAssistant - -from ..helpers.const import DOMAIN_SENSOR -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - - -class CoreSensor(SensorEntity, BaseEntity): - """Representation a binary sensor.""" - - @property - def native_value(self): - """Return the state of the sensor.""" - return self.entity.state - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - sensor = CoreSensor() - sensor.initialize(hass, entity, DOMAIN_SENSOR) - - return sensor - - @staticmethod - def get_domain(): - return DOMAIN_SENSOR diff --git a/custom_components/edgeos/core/components/switch.py b/custom_components/edgeos/core/components/switch.py deleted file mode 100644 index 21b57e9..0000000 --- a/custom_components/edgeos/core/components/switch.py +++ /dev/null @@ -1,48 +0,0 @@ -""" -Support for switch. -""" -from __future__ import annotations - -from typing import Any - -from homeassistant.components.switch import SwitchEntity -from homeassistant.core import HomeAssistant - -from ..helpers.const import DOMAIN_SWITCH -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - - -class CoreSwitch(SwitchEntity, BaseEntity): - """Class for a switch.""" - - @property - def is_on(self) -> bool | None: - """Return the boolean response if the node is on.""" - return self.entity.state - - def turn_on(self, **kwargs: Any) -> None: - self.hass.async_create_task(self.async_turn_on()) - - async def async_turn_on(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_turn_on(self.entity) - - def turn_off(self, **kwargs: Any) -> None: - self.hass.async_create_task(self.async_turn_off()) - - async def async_turn_off(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_turn_off(self.entity) - - async def async_setup(self): - pass - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - switch = CoreSwitch() - switch.initialize(hass, entity, DOMAIN_SWITCH) - - return switch - - @staticmethod - def get_domain(): - return DOMAIN_SWITCH diff --git a/custom_components/edgeos/core/components/vacuum.py b/custom_components/edgeos/core/components/vacuum.py deleted file mode 100644 index 262541e..0000000 --- a/custom_components/edgeos/core/components/vacuum.py +++ /dev/null @@ -1,105 +0,0 @@ -from __future__ import annotations - -from abc import ABC -import logging -import sys -from typing import Any - -from homeassistant.components.vacuum import StateVacuumEntity -from homeassistant.core import HomeAssistant - -from ..helpers.const import ATTR_FANS_SPEED_LIST, ATTR_FEATURES, DOMAIN_VACUUM -from ..models.base_entity import BaseEntity -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class CoreVacuum(StateVacuumEntity, BaseEntity, ABC): - """Class for a switch.""" - - def initialize( - self, - hass: HomeAssistant, - entity: EntityData, - current_domain: str, - ): - super().initialize(hass, entity, current_domain) - - try: - if hasattr(self.entity_description, ATTR_FEATURES): - self._attr_supported_features = getattr( - self.entity_description, ATTR_FEATURES - ) - - if hasattr(self.entity_description, ATTR_FANS_SPEED_LIST): - self._attr_fan_speed_list = getattr( - self.entity_description, ATTR_FANS_SPEED_LIST - ) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to initialize CoreSelect instance, Error: {ex}, Line: {line_number}" - ) - - @property - def state(self) -> str | None: - """Return the status of the vacuum cleaner.""" - return self.entity.state - - @property - def fan_speed(self) -> str | None: - """Return the fan speed of the vacuum cleaner.""" - return self.ha.get_core_entity_fan_speed(self.entity) - - async def async_return_to_base(self, **kwargs: Any) -> None: - """Set the vacuum cleaner to return to the dock.""" - await self.ha.async_core_entity_return_to_base(self.entity) - - async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None: - await self.ha.async_core_entity_set_fan_speed(self.entity, fan_speed) - - async def async_start(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_start(self.entity) - - async def async_stop(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_stop(self.entity) - - async def async_pause(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_pause(self.entity) - - async def async_turn_on(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_turn_on(self.entity) - - async def async_turn_off(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_turn_off(self.entity) - - async def async_toggle(self, **kwargs: Any) -> None: - await self.ha.async_core_entity_toggle(self.entity) - - async def async_send_command( - self, - command: str, - params: dict[str, Any] | list[Any] | None = None, - **kwargs: Any, - ) -> None: - """Send a command to a vacuum cleaner.""" - await self.ha.async_core_entity_send_command(self.entity, command, params) - - async def async_locate(self, **kwargs: Any) -> None: - """Locate the vacuum cleaner.""" - await self.ha.async_core_entity_locate(self.entity) - - @staticmethod - def get_component(hass: HomeAssistant, entity: EntityData): - vacuum = CoreVacuum() - vacuum.initialize(hass, entity, DOMAIN_VACUUM) - - return vacuum - - @staticmethod - def get_domain(): - return DOMAIN_VACUUM diff --git a/custom_components/edgeos/core/helpers/__init__.py b/custom_components/edgeos/core/helpers/__init__.py deleted file mode 100644 index 13e4d8c..0000000 --- a/custom_components/edgeos/core/helpers/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -from homeassistant.core import HomeAssistant - -from ...core.helpers.const import DATA - - -def get_ha(hass: HomeAssistant, entry_id): - ha_data = hass.data.get(DATA, {}) - ha = ha_data.get(entry_id) - - return ha diff --git a/custom_components/edgeos/core/helpers/const.py b/custom_components/edgeos/core/helpers/const.py deleted file mode 100644 index a0b4616..0000000 --- a/custom_components/edgeos/core/helpers/const.py +++ /dev/null @@ -1,76 +0,0 @@ -from homeassistant.components.binary_sensor import DOMAIN as DOMAIN_BINARY_SENSOR -from homeassistant.components.camera import DOMAIN as DOMAIN_CAMERA -from homeassistant.components.device_tracker import DOMAIN as DOMAIN_DEVICE_TRACKER -from homeassistant.components.light import DOMAIN as DOMAIN_LIGHT -from homeassistant.components.select import DOMAIN as DOMAIN_SELECT -from homeassistant.components.sensor import DOMAIN as DOMAIN_SENSOR -from homeassistant.components.switch import DOMAIN as DOMAIN_SWITCH -from homeassistant.components.vacuum import DOMAIN as DOMAIN_VACUUM - -from ...configuration.helpers.const import DOMAIN - -SUPPORTED_PLATFORMS = [ - DOMAIN_BINARY_SENSOR, - DOMAIN_CAMERA, - DOMAIN_SELECT, - DOMAIN_SWITCH, - DOMAIN_VACUUM, - DOMAIN_SENSOR, - DOMAIN_LIGHT, - DOMAIN_DEVICE_TRACKER, -] - -PLATFORMS = { - domain: f"{DOMAIN}_{domain}_UPDATE_SIGNAL" for domain in SUPPORTED_PLATFORMS -} - -ENTITY_STATE = "state" -ENTITY_ATTRIBUTES = "attributes" -ENTITY_UNIQUE_ID = "unique-id" -ENTITY_DEVICE_NAME = "device-name" -ENTITY_DETAILS = "details" -ENTITY_DISABLED = "disabled" -ENTITY_DOMAIN = "domain" -ENTITY_STATUS = "status" -ENTITY_CONFIG_ENTRY_ID = "entry_id" - -HA_NAME = "homeassistant" -SERVICE_RELOAD = "reload_config_entry" - -STORAGE_VERSION = 1 - -PASSWORD_MANAGER = f"pm_{DOMAIN}" -DATA = f"data_{DOMAIN}" - -DOMAIN_KEY_FILE = f"{DOMAIN}.key" - -ATTR_OPTIONS = "options" - -CONF_STILL_IMAGE_URL = "still_image_url" -CONF_STREAM_SOURCE = "stream_source" -CONF_MOTION_DETECTION = "motion_detection" - -ATTR_STREAM_FPS = "stream_fps" -ATTR_MODE_RECORD = "record_mode" -ATTR_FEATURES = "features" -ATTR_FANS_SPEED_LIST = "fan_speed_list" - -PROTOCOLS = {True: "https", False: "http"} -WS_PROTOCOLS = {True: "wss", False: "ws"} - -ACTION_CORE_ENTITY_RETURN_TO_BASE = "return_to_base" -ACTION_CORE_ENTITY_SET_FAN_SPEED = "set_fan_speed" -ACTION_CORE_ENTITY_START = "start" -ACTION_CORE_ENTITY_STOP = "stop" -ACTION_CORE_ENTITY_PAUSE = "stop" -ACTION_CORE_ENTITY_TURN_ON = "turn_on" -ACTION_CORE_ENTITY_TURN_OFF = "turn_off" -ACTION_CORE_ENTITY_TOGGLE = "toggle" -ACTION_CORE_ENTITY_SEND_COMMAND = "send_command" -ACTION_CORE_ENTITY_LOCATE = "locate" -ACTION_CORE_ENTITY_SELECT_OPTION = "select_option" -ACTION_CORE_ENTITY_ENABLE_MOTION_DETECTION = "enable_motion_detection" -ACTION_CORE_ENTITY_DISABLE_MOTION_DETECTION = "disable_motion_detection" - -SINGLE_FRAME_PS = 1 -EMPTY_STRING = "" diff --git a/custom_components/edgeos/core/helpers/enums.py b/custom_components/edgeos/core/helpers/enums.py deleted file mode 100644 index ba787db..0000000 --- a/custom_components/edgeos/core/helpers/enums.py +++ /dev/null @@ -1,34 +0,0 @@ -import logging - -from homeassistant.backports.enum import StrEnum - - -class ConnectivityStatus(StrEnum): - NotConnected = "Not connected" - Connecting = "Establishing connection to API" - Connected = "Connected to the API" - TemporaryConnected = "Connected with temporary API key" - Failed = "Failed to access API" - InvalidCredentials = "Invalid credentials" - MissingAPIKey = "Permanent API Key was not found" - Disconnected = "Disconnected by the system" - NotFound = "API Not found" - - @staticmethod - def get_log_level(status: StrEnum) -> int: - if status == ConnectivityStatus.Connected: - return logging.DEBUG - elif status in [ConnectivityStatus.Disconnected]: - return logging.INFO - elif status in [ConnectivityStatus.NotConnected, ConnectivityStatus.Connecting]: - return logging.WARNING - else: - return logging.ERROR - - -class EntityStatus(StrEnum): - EMPTY = "empty" - READY = "ready" - CREATED = "created" - DELETED = "deleted" - UPDATED = "updated" diff --git a/custom_components/edgeos/core/helpers/setup_base_entry.py b/custom_components/edgeos/core/helpers/setup_base_entry.py deleted file mode 100644 index 3c4f073..0000000 --- a/custom_components/edgeos/core/helpers/setup_base_entry.py +++ /dev/null @@ -1,43 +0,0 @@ -from collections.abc import Callable -import logging -import sys -from typing import Any - -from homeassistant.config_entries import ConfigEntry -from homeassistant.core import HomeAssistant -from homeassistant.helpers.entity_platform import AddEntitiesCallback - -from ..helpers.const import DATA -from ..models.domain_data import DomainData -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -async def async_setup_base_entry( - hass: HomeAssistant, - entry: ConfigEntry, - async_add_devices: AddEntitiesCallback, - domain: str, - initializer: Callable[[HomeAssistant, EntityData], Any], -): - """Set up base entity an entry.""" - _LOGGER.debug(f"Starting async_setup_entry {domain}") - - try: - ha_data = hass.data.get(DATA, {}) - - ha = ha_data.get(entry.entry_id) - - entity_manager = ha.entity_manager - - domain_data = DomainData(domain, async_add_devices, initializer) - - _LOGGER.debug(f"{domain} domain data: {domain_data}") - - entity_manager.set_domain_data(domain_data) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error(f"Failed to load {domain}, error: {ex}, line: {line_number}") diff --git a/custom_components/edgeos/core/managers/__init__.py b/custom_components/edgeos/core/managers/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/core/managers/device_manager.py b/custom_components/edgeos/core/managers/device_manager.py deleted file mode 100644 index f1c2934..0000000 --- a/custom_components/edgeos/core/managers/device_manager.py +++ /dev/null @@ -1,52 +0,0 @@ -import logging - -from homeassistant.helpers.device_registry import async_get - -_LOGGER = logging.getLogger(__name__) - - -class DeviceManager: - def __init__(self, hass, ha): - self._hass = hass - self._ha = ha - - self._devices = {} - - @property - def title(self): - title = self._ha.config_data.entry.title - - return title - - @property - def devices(self): - return self._devices - - async def async_remove_entry(self, entry_id): - dr = async_get(self._hass) - dr.async_clear_config_entry(entry_id) - - async def delete_device(self, name): - _LOGGER.info(f"Deleting device {name}") - - device = self._devices[name] - - device_identifiers = device.get("identifiers") - device_connections = device.get("connections", {}) - - dr = async_get(self._hass) - - device = dr.async_get_device(device_identifiers, device_connections) - - if device is not None: - dr.async_remove_device(device.id) - - async def async_remove(self): - for device_name in self._devices: - await self.delete_device(device_name) - - def get(self, name): - return self._devices.get(name, {}) - - def set(self, name, device_info): - self._devices[name] = device_info diff --git a/custom_components/edgeos/core/managers/entity_manager.py b/custom_components/edgeos/core/managers/entity_manager.py deleted file mode 100644 index e4b1886..0000000 --- a/custom_components/edgeos/core/managers/entity_manager.py +++ /dev/null @@ -1,311 +0,0 @@ -from __future__ import annotations - -from datetime import datetime -import json -import logging -import sys - -from homeassistant.core import HomeAssistant -from homeassistant.helpers.entity import EntityDescription -from homeassistant.helpers.entity_registry import EntityRegistry, RegistryEntryDisabler - -from ..helpers.const import DOMAIN -from ..helpers.enums import EntityStatus -from ..models.domain_data import DomainData -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class EntityManager: - """Entity Manager is agnostic to component - PLEASE DON'T CHANGE""" - - hass: HomeAssistant - domain_component_manager: dict[str, DomainData] - entities: dict[str, EntityData] - - def __init__(self, hass, ha): - self.hass: HomeAssistant = hass - self._ha = ha - self.domain_component_manager: dict[str, DomainData] = {} - self.entities = {} - - @property - def entity_registry(self) -> EntityRegistry: - return self._ha.entity_registry - - @property - def available_domains(self): - return self.domain_component_manager.keys() - - def set_domain_data(self, domain_data: DomainData): - self.domain_component_manager[domain_data.name] = domain_data - - def get_domain_data(self, domain: str) -> DomainData | None: - domain_data = self.domain_component_manager[domain] - - return domain_data - - def update(self): - self.hass.async_create_task(self._async_update()) - - async def _handle_disabled_entity(self, entity_id, entity: EntityData): - entity_item = self.entity_registry.async_get(entity_id) - - if entity_item is not None: - if entity.disabled: - _LOGGER.info(f"Disabling entity, Data: {entity}") - - self.entity_registry.async_update_entity( - entity_id, disabled_by=RegistryEntryDisabler.INTEGRATION - ) - - else: - entity.disabled = entity_item.disabled - - async def _handle_restored_entity(self, entity_id, component): - if entity_id is not None: - component.entity_id = entity_id - state = self.hass.states.get(entity_id) - - if state is not None: - restored = state.attributes.get("restored", False) - - if restored: - _LOGGER.debug(f"Restored {entity_id} ({component.name})") - - async def _async_add_components(self): - try: - components: dict[str, list] = {} - for unique_id in self.entities: - entity = self.entities.get(unique_id) - domain_manager = self.domain_component_manager.get(entity.domain) - - if entity.status == EntityStatus.CREATED and domain_manager is not None: - entity_id = self.entity_registry.async_get_entity_id( - entity.domain, DOMAIN, unique_id - ) - - await self._handle_disabled_entity(entity_id, entity) - - component = domain_manager.initializer(self.hass, entity) - - await self._handle_restored_entity(entity_id, component) - - domain_components = components.get(entity.domain, []) - domain_components.append(component) - - components[entity.domain] = domain_components - - entity.status = EntityStatus.READY - - elif entity.status == EntityStatus.UPDATED: - entity.status = EntityStatus.READY - - for domain in self.domain_component_manager: - domain_manager = self.domain_component_manager.get(domain) - - domain_components = components.get(domain, []) - components_count = len(domain_components) - - if components_count > 0: - domain_manager.async_add_devices(domain_components) - - _LOGGER.info(f"{components_count} {domain} components created") - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to add component, " - f"Error: {str(ex)}, " - f"Line: {line_number}" - ) - - async def _async_delete_components(self): - try: - delete_entities = [] - for unique_id in self.entities: - entity = self.entities.get(unique_id) - - if entity.status == EntityStatus.DELETED: - entity_id = self.entity_registry.async_get_entity_id( - entity.domain, DOMAIN, unique_id - ) - - entity_item = self.entity_registry.async_get(entity_id) - - if entity_item is not None: - _LOGGER.info(f"Removed {entity_id} ({entity.name})") - - self.entity_registry.async_remove(entity_id) - - delete_entities.append(unique_id) - - for unique_id in delete_entities: - self.entities.pop(unique_id, None) - - total_delete_entities = len(delete_entities) - if total_delete_entities > 0: - _LOGGER.info(f"{total_delete_entities} components deleted") - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to delete components, " - f"Error: {str(ex)}, " - f"Line: {line_number}" - ) - - async def _async_update(self): - try: - await self._async_add_components() - await self._async_delete_components() - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to update, " f"Error: {str(ex)}, " f"Line: {line_number}" - ) - - def _compare_data( - self, - entity_name: str, - entity: EntityData, - state: str | int | float | bool, - attributes: dict, - device_name: str, - entity_description: EntityDescription | None = None, - details: dict | None = None, - ): - msgs = [] - - if str(entity.state) != str(state): - msgs.append(f"State {entity.state} -> {state}") - - if entity.attributes != attributes: - from_attributes = self._get_attributes_json(entity.attributes) - to_attributes = self._get_attributes_json(attributes) - - msgs.append(f"Attributes {from_attributes} -> {to_attributes}") - - if entity.device_name != device_name: - msgs.append(f"Device name {entity.device_name} -> {device_name}") - - if ( - entity_description is not None - and entity.entity_description != entity_description - ): - msgs.append( - f"Description {str(entity.entity_description)} -> {str(entity_description)}" - ) - - if details is not None and entity.details != details: - from_details = self._get_attributes_json(entity.details) - to_details = self._get_attributes_json(details) - - msgs.append(f"Details {from_details} -> {to_details}") - - modified = len(msgs) > 0 - - if modified: - full_message = " | ".join(msgs) - - _LOGGER.debug(f"{entity_name} | {entity.domain} | {full_message}") - - return modified - - @staticmethod - def _get_attributes_json(attributes: dict): - new_attributes = {} - for key in attributes: - value = attributes[key] - new_attributes[key] = str(value) - - result = json.dumps(new_attributes) - - return result - - def get(self, unique_id: str) -> EntityData | None: - entity = self.entities.get(unique_id) - - return entity - - def set_entity( - self, - domain: str, - entry_id: str, - state: str | int | float | bool | datetime, - attributes: dict, - device_name: str, - entity_description: EntityDescription | None, - details: dict | None = None, - destructors: list[bool] = None, - ): - try: - entity = self.entities.get(entity_description.key) - entity_name = entity_description.name - original_status = None - - if destructors is not None and True in destructors: - if entity is not None and entity.status != EntityStatus.CREATED: - _LOGGER.debug(f"{entity_name} will be removed") - - entity.status = EntityStatus.DELETED - - self.entities[entity_description.key] = entity - - else: - if entity is None: - entity = EntityData(entry_id) - entity.status = EntityStatus.CREATED - entity.domain = domain - - self._compare_data( - entity_name, entity, state, attributes, device_name - ) - - else: - original_status = entity.status - was_modified = self._compare_data( - entity_name, - entity, - state, - attributes, - device_name, - entity_description, - details, - ) - - if was_modified: - entity.status = EntityStatus.UPDATED - - if entity.status in [EntityStatus.CREATED, EntityStatus.UPDATED]: - entity.state = state - entity.attributes = attributes - entity.device_name = device_name - entity.details = details - entity.entity_description = entity_description - - self.entities[entity_description.key] = entity - - if entity.status != EntityStatus.READY: - _LOGGER.info( - f"{entity_name} ({entity.domain}) {original_status} -> {entity.status}, " - f"state: {entity.state}" - ) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to set entity {entity_description.name}, " - f"Error: {str(ex)}, " - f"Line: {line_number}" - ) diff --git a/custom_components/edgeos/core/managers/home_assistant.py b/custom_components/edgeos/core/managers/home_assistant.py deleted file mode 100644 index 95555c4..0000000 --- a/custom_components/edgeos/core/managers/home_assistant.py +++ /dev/null @@ -1,423 +0,0 @@ -""" -Core HA Manager. -""" -from __future__ import annotations - -import datetime -import logging -import sys -from typing import Any - -from cryptography.fernet import InvalidToken - -from homeassistant.config_entries import ConfigEntry -from homeassistant.core import HomeAssistant -from homeassistant.helpers.dispatcher import async_dispatcher_send -from homeassistant.helpers.entity_registry import EntityRegistry, async_get -from homeassistant.helpers.event import async_track_time_interval - -from ..helpers.const import ( - ACTION_CORE_ENTITY_DISABLE_MOTION_DETECTION, - ACTION_CORE_ENTITY_ENABLE_MOTION_DETECTION, - ACTION_CORE_ENTITY_LOCATE, - ACTION_CORE_ENTITY_PAUSE, - ACTION_CORE_ENTITY_RETURN_TO_BASE, - ACTION_CORE_ENTITY_SELECT_OPTION, - ACTION_CORE_ENTITY_SEND_COMMAND, - ACTION_CORE_ENTITY_SET_FAN_SPEED, - ACTION_CORE_ENTITY_START, - ACTION_CORE_ENTITY_STOP, - ACTION_CORE_ENTITY_TOGGLE, - ACTION_CORE_ENTITY_TURN_OFF, - ACTION_CORE_ENTITY_TURN_ON, - DOMAIN, - PLATFORMS, - SUPPORTED_PLATFORMS, -) -from ..managers.device_manager import DeviceManager -from ..managers.entity_manager import EntityManager -from ..managers.storage_manager import StorageManager -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class HomeAssistantManager: - def __init__( - self, - hass: HomeAssistant, - scan_interval: datetime.timedelta, - heartbeat_interval: datetime.timedelta | None = None, - ): - self._hass = hass - - self._is_initialized = False - self._update_entities_interval = scan_interval - self._update_data_providers_interval = scan_interval - self._heartbeat_interval = heartbeat_interval - - self._entity_registry = None - - self._entry: ConfigEntry | None = None - - self._storage_manager = StorageManager(self._hass) - self._entity_manager = EntityManager(self._hass, self) - self._device_manager = DeviceManager(self._hass, self) - - self._entity_registry = async_get(self._hass) - - self._async_track_time_handlers = [] - self._last_heartbeat = None - self._update_lock = False - self._actions: dict = {} - - def _send_heartbeat(internal_now): - self._last_heartbeat = internal_now - - self._hass.async_create_task(self.async_send_heartbeat()) - - self._send_heartbeat = _send_heartbeat - - self._domains = { - domain: self.is_domain_supported(domain) for domain in SUPPORTED_PLATFORMS - } - - @property - def entity_manager(self) -> EntityManager: - if self._entity_manager is None: - self._entity_manager = EntityManager(self._hass, self) - - return self._entity_manager - - @property - def device_manager(self) -> DeviceManager: - return self._device_manager - - @property - def entity_registry(self) -> EntityRegistry: - return self._entity_registry - - @property - def storage_manager(self) -> StorageManager: - return self._storage_manager - - @property - def entry_id(self) -> str: - return self._entry.entry_id - - @property - def entry_title(self) -> str: - return self._entry.title - - def update_intervals( - self, entities_interval: datetime.timedelta, data_interval: datetime.timedelta - ): - self._update_entities_interval = entities_interval - self._update_data_providers_interval = data_interval - - async def async_component_initialize(self, entry: ConfigEntry): - """Component initialization""" - - async def async_send_heartbeat(self): - """Must be implemented to be able to send heartbeat to API""" - - def register_services(self, entry: ConfigEntry | None = None): - """Must be implemented to be able to expose services""" - - async def async_initialize_data_providers(self): - """Must be implemented to be able to send heartbeat to API""" - - async def async_stop_data_providers(self): - """Must be implemented to be able to send heartbeat to API""" - - async def async_update_data_providers(self): - """Must be implemented to be able to send heartbeat to API""" - - def load_entities(self): - """Must be implemented to be able to send heartbeat to API""" - - def load_devices(self): - """Must be implemented to be able to send heartbeat to API""" - - async def async_init(self, entry: ConfigEntry): - try: - self._entry = entry - - await self.async_component_initialize(entry) - - self._hass.loop.create_task(self._async_load_platforms()) - - except InvalidToken: - error_message = "Encryption key got corrupted, please remove the integration and re-add it" - - _LOGGER.error(error_message) - - data = await self._storage_manager.async_load_from_store() - data.key = None - - await self._storage_manager.async_save_to_store(data) - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error(f"Failed to async_init, error: {ex}, line: {line_number}") - - async def _async_load_platforms(self): - load = self._hass.config_entries.async_forward_entry_setup - - for domain in self._domains: - if self._domains.get(domain, False): - await load(self._entry, domain) - - else: - _LOGGER.debug(f"Skip loading {domain}") - - self.register_services() - - self._is_initialized = True - - await self.async_update_entry() - - def _update_data_providers(self, now): - self._hass.async_create_task(self.async_update_data_providers()) - - async def async_update_entry(self, entry: ConfigEntry | None = None): - entry_changed = entry is not None - - if entry_changed: - self._entry = entry - - _LOGGER.info(f"Handling ConfigEntry load: {entry.as_dict()}") - - else: - entry = self._entry - - track_time_update_data_providers = async_track_time_interval( - self._hass, - self._update_data_providers, - self._update_data_providers_interval, - ) - - self._async_track_time_handlers.append(track_time_update_data_providers) - - track_time_update_entities = async_track_time_interval( - self._hass, self._update_entities, self._update_entities_interval - ) - - self._async_track_time_handlers.append(track_time_update_entities) - - if self._heartbeat_interval is not None: - track_time_send_heartbeat = async_track_time_interval( - self._hass, self._send_heartbeat, self._heartbeat_interval - ) - - self._async_track_time_handlers.append(track_time_send_heartbeat) - - _LOGGER.info(f"Handling ConfigEntry change: {entry.as_dict()}") - - await self.async_initialize_data_providers() - - async def async_unload(self): - _LOGGER.info("HA was stopped") - - for handler in self._async_track_time_handlers: - if handler is not None: - handler() - - self._async_track_time_handlers.clear() - - await self.async_stop_data_providers() - - async def async_remove(self, entry: ConfigEntry): - _LOGGER.info(f"Removing current integration - {entry.title}") - - await self.async_unload() - - unload = self._hass.config_entries.async_forward_entry_unload - - for domain in PLATFORMS: - if self._domains.get(domain, False): - await unload(self._entry, domain) - - else: - _LOGGER.debug(f"Skip unloading {domain}") - - await self._device_manager.async_remove() - - self._entry = None - self.entity_manager.entities.clear() - - _LOGGER.info(f"Current integration ({entry.title}) removed") - - def _update_entities(self, now): - if self._update_lock: - _LOGGER.warning("Update in progress, will skip the request") - return - - self._update_lock = True - - try: - self.load_devices() - self.load_entities() - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to update devices and entities, Error: {ex}, Line: {line_number}" - ) - - self.entity_manager.update() - - self._hass.async_create_task(self.dispatch_all()) - - self._update_lock = False - - async def dispatch_all(self): - if not self._is_initialized: - _LOGGER.info("NOT INITIALIZED - Failed discovering components") - return - - for domain in PLATFORMS: - if self._domains.get(domain, False): - signal = PLATFORMS.get(domain) - - async_dispatcher_send(self._hass, signal) - - def set_action(self, entity_id: str, action_name: str, action): - key = f"{entity_id}:{action_name}" - self._actions[key] = action - - def get_action(self, entity_id: str, action_name: str): - key = f"{entity_id}:{action_name}" - action = self._actions.get(key) - - return action - - def get_core_entity_fan_speed(self, entity: EntityData) -> str | None: - pass - - async def async_core_entity_return_to_base(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_RETURN_TO_BASE.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_RETURN_TO_BASE) - - if action is not None: - await action(entity) - - async def async_core_entity_set_fan_speed( - self, entity: EntityData, fan_speed: str - ) -> None: - """Handles ACTION_CORE_ENTITY_SET_FAN_SPEED.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_SET_FAN_SPEED) - - if action is not None: - await action(entity, fan_speed) - - async def async_core_entity_start(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_START.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_START) - - if action is not None: - await action(entity) - - async def async_core_entity_stop(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_STOP.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_STOP) - - if action is not None: - await action(entity) - - async def async_core_entity_pause(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_PAUSE.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_PAUSE) - - if action is not None: - await action(entity) - - async def async_core_entity_turn_on(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_TURN_ON.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_TURN_ON) - - if action is not None: - await action(entity) - - async def async_core_entity_turn_off(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_TURN_OFF.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_TURN_OFF) - - if action is not None: - await action(entity) - - async def async_core_entity_send_command( - self, - entity: EntityData, - command: str, - params: dict[str, Any] | list[Any] | None = None, - ) -> None: - """Handles ACTION_CORE_ENTITY_SEND_COMMAND.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_SEND_COMMAND) - - if action is not None: - await action(entity, command, params) - - async def async_core_entity_locate(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_LOCATE.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_LOCATE) - - if action is not None: - await action(entity) - - async def async_core_entity_select_option( - self, entity: EntityData, option: str - ) -> None: - """Handles ACTION_CORE_ENTITY_SELECT_OPTION.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_SELECT_OPTION) - - if action is not None: - await action(entity, option) - - async def async_core_entity_toggle(self, entity: EntityData) -> None: - """Handles ACTION_CORE_ENTITY_TOGGLE.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_TOGGLE) - - if action is not None: - await action(entity) - - async def async_core_entity_enable_motion_detection( - self, entity: EntityData - ) -> None: - """Handles ACTION_CORE_ENTITY_ENABLE_MOTION_DETECTION.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_ENABLE_MOTION_DETECTION) - - if action is not None: - await action(entity) - - async def async_core_entity_disable_motion_detection( - self, entity: EntityData - ) -> None: - """Handles ACTION_CORE_ENTITY_DISABLE_MOTION_DETECTION.""" - action = self.get_action(entity.id, ACTION_CORE_ENTITY_DISABLE_MOTION_DETECTION) - - if action is not None: - await action(entity) - - @staticmethod - def log_exception(ex, message): - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error(f"{message}, Error: {str(ex)}, Line: {line_number}") - - @staticmethod - def is_domain_supported(domain) -> bool: - is_supported = True - - try: - __import__(f"custom_components.{DOMAIN}.{domain}") - - except ModuleNotFoundError: - is_supported = False - - return is_supported diff --git a/custom_components/edgeos/core/managers/password_manager.py b/custom_components/edgeos/core/managers/password_manager.py deleted file mode 100644 index 3466679..0000000 --- a/custom_components/edgeos/core/managers/password_manager.py +++ /dev/null @@ -1,66 +0,0 @@ -from __future__ import annotations - -import logging -from os import path, remove -import sys - -from cryptography.fernet import Fernet - -from homeassistant.core import HomeAssistant - -from ..helpers.const import DOMAIN_KEY_FILE -from ..managers.storage_manager import StorageManager -from ..models.storage_data import StorageData - -_LOGGER = logging.getLogger(__name__) - - -class PasswordManager: - data: StorageData | None - hass: HomeAssistant - crypto: Fernet - - def __init__(self, hass: HomeAssistant): - self.hass = hass - self.data = None - - async def initialize(self): - try: - if self.data is None: - storage_manager = StorageManager(self.hass) - - self.data = await storage_manager.async_load_from_store() - - if self.data.key is None: - legacy_key_path = self.hass.config.path(DOMAIN_KEY_FILE) - - if path.exists(legacy_key_path): - with open(legacy_key_path, "rb") as file: - self.data.key = file.read().decode("utf-8") - - remove(legacy_key_path) - else: - self.data.key = Fernet.generate_key().decode("utf-8") - - await storage_manager.async_save_to_store(self.data) - - self.crypto = Fernet(self.data.key.encode()) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to initialize Password Manager, error: {ex}, line: {line_number}" - ) - - def set(self, data: str) -> str: - if data is not None: - data = self.crypto.encrypt(data.encode()).decode() - - return data - - def get(self, data: str) -> str: - if data is not None and len(data) > 0: - data = self.crypto.decrypt(data.encode()).decode() - - return data diff --git a/custom_components/edgeos/core/managers/storage_manager.py b/custom_components/edgeos/core/managers/storage_manager.py deleted file mode 100644 index d0ac584..0000000 --- a/custom_components/edgeos/core/managers/storage_manager.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Storage handlers.""" -import logging - -from homeassistant.helpers.json import JSONEncoder -from homeassistant.helpers.storage import Store - -from ..helpers.const import DOMAIN, STORAGE_VERSION -from ..models.storage_data import StorageData - -_LOGGER = logging.getLogger(__name__) - - -class StorageManager: - def __init__(self, hass): - self._hass = hass - - @property - def file_name(self): - file_name = f".{DOMAIN}" - - return file_name - - async def async_load_from_store(self) -> StorageData: - """Load the retained data from store and return de-serialized data.""" - store = Store(self._hass, STORAGE_VERSION, self.file_name, encoder=JSONEncoder) - - data = await store.async_load() - - result = StorageData.from_dict(data) - - return result - - async def async_save_to_store(self, data: StorageData): - """Generate dynamic data to store and save it to the filesystem.""" - store = Store(self._hass, STORAGE_VERSION, self.file_name, encoder=JSONEncoder) - - await store.async_save(data.to_dict()) diff --git a/custom_components/edgeos/core/models/__init__.py b/custom_components/edgeos/core/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/custom_components/edgeos/core/models/base_entity.py b/custom_components/edgeos/core/models/base_entity.py deleted file mode 100644 index bc7b12e..0000000 --- a/custom_components/edgeos/core/models/base_entity.py +++ /dev/null @@ -1,133 +0,0 @@ -from __future__ import annotations - -import logging -import sys - -from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.helpers.entity import Entity - -from ..helpers.const import DATA, PLATFORMS -from ..managers.device_manager import DeviceManager -from ..managers.entity_manager import EntityManager -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class BaseEntity(Entity): - """Representation a base entity.""" - - hass: HomeAssistant | None = None - entity: EntityData | None = None - remove_dispatcher = None - current_domain: str = None - - ha = None - entity_manager: EntityManager = None - device_manager: DeviceManager = None - - def initialize( - self, - hass: HomeAssistant, - entity: EntityData, - current_domain: str, - ): - try: - self.hass = hass - self.entity = entity - self.remove_dispatcher = None - self.current_domain = current_domain - - ha_data = hass.data.get(DATA, {}) - - self.ha = ha_data.get(entity.entry_id) - - if self.ha is None: - _LOGGER.warning("Failed to initialize BaseEntity without HA manager") - return - - self.entity_manager = self.ha.entity_manager - self.device_manager = self.ha.device_manager - self.entity_description = entity.entity_description - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to initialize BaseEntity, Error: {ex}, Line: {line_number}" - ) - - @property - def entry_id(self) -> str | None: - """Return the name of the node.""" - return self.entity.entry_id - - @property - def unique_id(self) -> str | None: - """Return the name of the node.""" - return self.entity.id - - @property - def device_info(self): - return self.device_manager.get(self.entity.device_name) - - @property - def name(self): - """Return the name of the node.""" - return self.entity.name - - @property - def extra_state_attributes(self): - """Return true if the binary sensor is on.""" - return self.entity.attributes - - async def async_added_to_hass(self): - """Register callbacks.""" - async_dispatcher_connect( - self.hass, PLATFORMS[self.current_domain], self._schedule_immediate_update - ) - - await self.async_added_to_hass_local() - - async def async_will_remove_from_hass(self) -> None: - if self.remove_dispatcher is not None: - self.remove_dispatcher() - self.remove_dispatcher = None - - _LOGGER.debug(f"Removing component: {self.unique_id}") - - self.entity = None - - await self.async_will_remove_from_hass_local() - - @callback - def _schedule_immediate_update(self): - self.hass.async_create_task(self._async_schedule_immediate_update()) - - async def _async_schedule_immediate_update(self): - if self.entity_manager is None: - _LOGGER.debug( - f"Cannot update {self.current_domain} - Entity Manager is None | {self.name}" - ) - else: - if self.entity is not None: - entity = self.entity_manager.get(self.unique_id) - - if entity is None: - _LOGGER.debug(f"Skip updating {self.name}, Entity is None") - - elif entity.disabled: - _LOGGER.debug(f"Skip updating {self.name}, Entity is disabled") - - else: - self.entity = entity - if self.entity is not None: - self.async_schedule_update_ha_state(True) - - async def async_added_to_hass_local(self): - pass - - async def async_will_remove_from_hass_local(self): - pass diff --git a/custom_components/edgeos/core/models/domain_data.py b/custom_components/edgeos/core/models/domain_data.py deleted file mode 100644 index eec7444..0000000 --- a/custom_components/edgeos/core/models/domain_data.py +++ /dev/null @@ -1,38 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable -import logging -from typing import Any - -from homeassistant.const import CONF_NAME -from homeassistant.core import HomeAssistant -from homeassistant.helpers.entity_platform import AddEntitiesCallback - -from ..models.entity_data import EntityData - -_LOGGER = logging.getLogger(__name__) - - -class DomainData: - name: str - async_add_devices: AddEntitiesCallback - initializer: Callable[[HomeAssistant, EntityData], Any] - - def __init__( - self, - name, - async_add_devices: AddEntitiesCallback, - initializer: Callable[[HomeAssistant, EntityData], Any], - ): - self.name = name - self.async_add_devices = async_add_devices - self.initializer = initializer - - _LOGGER.info(f"Creating domain data for {name}") - - def __repr__(self): - obj = {CONF_NAME: self.name} - - to_string = f"{obj}" - - return to_string diff --git a/custom_components/edgeos/core/models/entity_data.py b/custom_components/edgeos/core/models/entity_data.py deleted file mode 100644 index 15f91df..0000000 --- a/custom_components/edgeos/core/models/entity_data.py +++ /dev/null @@ -1,73 +0,0 @@ -from __future__ import annotations - -from datetime import datetime - -from homeassistant.helpers.entity import EntityDescription -from homeassistant.util import slugify - -from ..helpers.const import ( - ENTITY_ATTRIBUTES, - ENTITY_CONFIG_ENTRY_ID, - ENTITY_DETAILS, - ENTITY_DEVICE_NAME, - ENTITY_DISABLED, - ENTITY_DOMAIN, - ENTITY_STATE, - ENTITY_STATUS, - ENTITY_UNIQUE_ID, -) -from ..helpers.enums import EntityStatus - - -class EntityData: - state: str | int | float | bool | datetime | None - attributes: dict - details: dict - device_name: str | None - status: EntityStatus - disabled: bool - domain: str | None - entry_id: str - entity_description: EntityDescription | None - - def __init__(self, entry_id: str): - self.entry_id = entry_id - self.entity_description = None - self.state = None - self.attributes = {} - self.details = {} - self.device_name = None - self.status = EntityStatus.CREATED - self.disabled = False - self.domain = None - - @property - def id(self): - return self.entity_description.key - - @property - def name(self): - return self.entity_description.name - - def __repr__(self): - obj = { - ENTITY_UNIQUE_ID: self.id, - ENTITY_STATE: self.state, - ENTITY_ATTRIBUTES: self.attributes, - ENTITY_DETAILS: self.details, - ENTITY_DEVICE_NAME: self.device_name, - ENTITY_STATUS: self.status, - ENTITY_DISABLED: self.disabled, - ENTITY_DOMAIN: self.domain, - ENTITY_CONFIG_ENTRY_ID: self.entry_id, - } - - to_string = f"{obj}" - - return to_string - - @staticmethod - def generate_unique_id(domain, name): - unique_id = slugify(f"{domain} {name}") - - return unique_id diff --git a/custom_components/edgeos/core/models/storage_data.py b/custom_components/edgeos/core/models/storage_data.py deleted file mode 100644 index c97c15b..0000000 --- a/custom_components/edgeos/core/models/storage_data.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - - -class StorageData: - key: str | None - - def __init__(self): - self.key = None - - @staticmethod - def from_dict(obj: dict): - data = StorageData() - - if obj is not None: - data.key = obj.get("key") - - return data - - def to_dict(self): - obj = {"key": self.key} - - return obj - - def __repr__(self): - to_string = f"{self.to_dict()}" - - return to_string diff --git a/custom_components/edgeos/core/models/vacuum_description.py b/custom_components/edgeos/core/models/vacuum_description.py deleted file mode 100644 index fac7e2d..0000000 --- a/custom_components/edgeos/core/models/vacuum_description.py +++ /dev/null @@ -1,11 +0,0 @@ -from dataclasses import dataclass - -from homeassistant.components.vacuum import StateVacuumEntityDescription - - -@dataclass(frozen=True, kw_only=True) -class VacuumDescription(StateVacuumEntityDescription): - """A class that describes vacuum entities.""" - - features: int = 0 - fan_speed_list: list[str] = () diff --git a/custom_components/edgeos/data_processors/base_processor.py b/custom_components/edgeos/data_processors/base_processor.py new file mode 100644 index 0000000..6493297 --- /dev/null +++ b/custom_components/edgeos/data_processors/base_processor.py @@ -0,0 +1,87 @@ +import logging + +from homeassistant.helpers.device_registry import DeviceInfo +from homeassistant.util import slugify + +from ..common.consts import ( + API_DATA_SYSTEM, + DATA_SYSTEM_SYSTEM, + DEFAULT_NAME, + SYSTEM_DATA_HOSTNAME, +) +from ..common.enums import DeviceTypes +from ..models.config_data import ConfigData + +_LOGGER = logging.getLogger(__name__) + + +class BaseProcessor: + _api_data: dict | None = None + _ws_data: dict | None = None + _config_data: ConfigData | None = None + _unique_messages: list[str] | None = None + processor_type: DeviceTypes | None = None + _hostname: str | None = None + + def __init__(self, config_data: ConfigData): + self._config_data = config_data + + self._api_data = None + self._ws_data = None + self.processor_type = None + self._hostname = None + + self._unique_messages = [] + + def update(self, api_data: dict, ws_data: dict): + self._api_data = api_data + self._ws_data = ws_data + + self._process_api_data() + self._process_ws_data() + + def _process_api_data(self): + system_section = self._api_data.get(API_DATA_SYSTEM, {}) + system_details = system_section.get(DATA_SYSTEM_SYSTEM, {}) + + self._hostname = system_details.get(SYSTEM_DATA_HOSTNAME).upper() + + def _process_ws_data(self): + pass + + def _unique_log(self, log_level: int, message: str): + if message not in self._unique_messages: + self._unique_messages.append(message) + + _LOGGER.log(log_level, message) + + def get_device_info(self, item_id: str | None = None) -> DeviceInfo: + device_name = self._get_device_info_name(item_id) + + unique_id = self._get_device_info_unique_id(item_id) + + device_info = DeviceInfo( + identifiers={(DEFAULT_NAME, unique_id)}, + name=device_name, + model=self.processor_type, + manufacturer=DEFAULT_NAME, + via_device=(DEFAULT_NAME, self._hostname), + ) + + return device_info + + def _get_device_info_name(self, item_id: str | None = None): + parts = [self._hostname, self.processor_type, item_id] + + relevant_parts = [part for part in parts if part is not None] + + name = " ".join(relevant_parts) + + return name + + def _get_device_info_unique_id(self, item_id: str | None = None): + identifier = self._get_device_info_name(item_id) + + unique_id = slugify(identifier) + + return unique_id diff --git a/custom_components/edgeos/data_processors/device_processor.py b/custom_components/edgeos/data_processors/device_processor.py new file mode 100644 index 0000000..3822bc2 --- /dev/null +++ b/custom_components/edgeos/data_processors/device_processor.py @@ -0,0 +1,245 @@ +import logging +import sys + +from homeassistant.helpers.device_registry import DeviceInfo + +from ..common.consts import ( + API_DATA_DHCP_LEASES, + API_DATA_SYSTEM, + DATA_SYSTEM_SERVICE, + DATA_SYSTEM_SERVICE_DHCP_SERVER, + DEFAULT_NAME, + DEVICE_DATA_MAC, + DHCP_SERVER_IP_ADDRESS, + DHCP_SERVER_LEASES, + DHCP_SERVER_LEASES_CLIENT_HOSTNAME, + DHCP_SERVER_MAC_ADDRESS, + DHCP_SERVER_SHARED_NETWORK_NAME, + DHCP_SERVER_STATIC_MAPPING, + DHCP_SERVER_SUBNET, + SYSTEM_DATA_DOMAIN_NAME, + TRAFFIC_DATA_DEVICE_ITEMS, + WS_EXPORT_KEY, +) +from ..common.enums import DeviceTypes +from ..models.config_data import ConfigData +from ..models.edge_os_device_data import EdgeOSDeviceData +from .base_processor import BaseProcessor + +_LOGGER = logging.getLogger(__name__) + + +class DeviceProcessor(BaseProcessor): + _devices: dict[str, EdgeOSDeviceData] + _devices_ip_mapping: dict[str, str] + + def __init__(self, config_data: ConfigData): + super().__init__(config_data) + + self.processor_type = DeviceTypes.DEVICE + + self._devices = {} + self._devices_ip_mapping = {} + self._leased_devices = {} + + def get_devices(self) -> list[str]: + return list(self._devices.keys()) + + def get_all(self) -> list[dict]: + items = [self._devices[item_key].to_dict() for item_key in self._devices] + + return items + + def get_device(self, identifiers: set[tuple[str, str]]) -> dict | None: + device: dict | None = None + device_identifier = list(identifiers)[0][1] + + for device_mac in self._devices: + unique_id = self._get_device_info_unique_id(device_mac) + + if unique_id == device_identifier: + device = self._devices[device_mac].to_dict() + + return device + + def get_data(self, device_mac: str) -> EdgeOSDeviceData: + interface_data = self._devices.get(device_mac) + + return interface_data + + def get_device_info(self, item_id: str | None = None) -> DeviceInfo: + device = self.get_data(item_id) + device_name = self._get_device_info_name(device.hostname) + + unique_id = self._get_device_info_unique_id(device.hostname) + + device_info = DeviceInfo( + identifiers={(DEFAULT_NAME, unique_id)}, + name=device_name, + model=self.processor_type, + manufacturer=DEFAULT_NAME, + via_device=(DEFAULT_NAME, self._hostname), + ) + + return device_info + + def get_leased_devices(self) -> dict: + return self._leased_devices + + def _process_api_data(self): + super()._process_api_data() + + try: + system_section = self._api_data.get(API_DATA_SYSTEM, {}) + service = system_section.get(DATA_SYSTEM_SERVICE, {}) + + dhcp_server = service.get(DATA_SYSTEM_SERVICE_DHCP_SERVER, {}) + shared_network_names = dhcp_server.get(DHCP_SERVER_SHARED_NETWORK_NAME, {}) + + for shared_network_name in shared_network_names: + shared_network_name_data = shared_network_names.get( + shared_network_name, {} + ) + subnets = shared_network_name_data.get(DHCP_SERVER_SUBNET, {}) + + for subnet in subnets: + subnet_data = subnets.get(subnet, {}) + + domain_name = subnet_data.get(SYSTEM_DATA_DOMAIN_NAME) + static_mappings = subnet_data.get(DHCP_SERVER_STATIC_MAPPING, {}) + + for hostname in static_mappings: + static_mapping_data = static_mappings.get(hostname, {}) + + self._set_device( + hostname, domain_name, static_mapping_data, False + ) + + self._update_leased_devices() + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract Devices data, Error: {ex}, Line: {line_number}" + ) + + def _process_ws_data(self): + try: + device_data = self._ws_data.get(WS_EXPORT_KEY, {}) + + for device_ip in device_data: + device_item = self._get_device_by_ip(device_ip) + stats = device_data.get(device_ip) + + if device_item is not None: + self._update_device_stats(device_item, stats) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract WS data, Error: {ex}, Line: {line_number}" + ) + + def _update_leased_devices(self): + try: + data_leases = self._api_data.get(API_DATA_DHCP_LEASES, {}) + data_server_leases = data_leases.get(DHCP_SERVER_LEASES, {}) + + for subnet in data_server_leases: + subnet_data = data_server_leases.get(subnet, {}) + + for ip in subnet_data: + device_data = subnet_data.get(ip) + + hostname = device_data.get(DHCP_SERVER_LEASES_CLIENT_HOSTNAME) + + static_mapping_data = { + DHCP_SERVER_IP_ADDRESS: ip, + DHCP_SERVER_MAC_ADDRESS: device_data.get(DEVICE_DATA_MAC), + } + + self._set_device(hostname, None, static_mapping_data, True) + + self._leased_devices.clear() + + for device_mac in self._devices: + device = self._devices.get(device_mac) + if device.is_leased: + device_name = device.mac + + if device.hostname not in ["", "?"]: + device_name = f"{device.mac} ({device.hostname})" + + self._leased_devices[device.ip] = device_name + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract Unknown Devices data, Error: {ex}, Line: {line_number}" + ) + + def _set_device( + self, + hostname: str, + domain_name: str | None, + static_mapping_data: dict, + is_leased: bool, + ): + ip_address = static_mapping_data.get(DHCP_SERVER_IP_ADDRESS) + mac_address = static_mapping_data.get(DHCP_SERVER_MAC_ADDRESS) + + existing_device_data = self._devices.get(mac_address) + + if existing_device_data is None: + device_data = EdgeOSDeviceData( + hostname, ip_address, mac_address, domain_name, is_leased + ) + + else: + device_data = existing_device_data + + self._devices[device_data.unique_id] = device_data + self._devices_ip_mapping[device_data.ip] = device_data.unique_id + + def _get_device(self, unique_id: str) -> EdgeOSDeviceData | None: + device = self._devices.get(unique_id) + + return device + + def _get_device_by_ip(self, ip: str) -> EdgeOSDeviceData | None: + unique_id = self._devices_ip_mapping.get(ip) + + device = self._get_device(unique_id) + + return device + + @staticmethod + def _update_device_stats(device_data: EdgeOSDeviceData, stats: dict): + try: + if not device_data.is_leased: + directions = [device_data.received, device_data.sent] + + for direction in directions: + stat_data = {} + for stat_key in TRAFFIC_DATA_DEVICE_ITEMS: + key = f"{direction.direction}_{stat_key}" + stat_data_item = TRAFFIC_DATA_DEVICE_ITEMS.get(stat_key) + + stat_data[stat_data_item] = stats.get(key) + + direction.update(stat_data) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to update device statistics for {device_data.hostname}, " + f"Error: {ex}, " + f"Line: {line_number}" + ) diff --git a/custom_components/edgeos/data_processors/interface_processor.py b/custom_components/edgeos/data_processors/interface_processor.py new file mode 100644 index 0000000..56a17a1 --- /dev/null +++ b/custom_components/edgeos/data_processors/interface_processor.py @@ -0,0 +1,223 @@ +import logging +import sys + +from homeassistant.helpers.device_registry import DeviceInfo + +from ..common.consts import ( + ADDRESS_LIST, + API_DATA_INTERFACES, + API_DATA_SYSTEM, + DEFAULT_NAME, + FALSE_STR, + INTERFACE_DATA_ADDRESS, + INTERFACE_DATA_AGING, + INTERFACE_DATA_BRIDGE_GROUP, + INTERFACE_DATA_BRIDGED_CONNTRACK, + INTERFACE_DATA_DESCRIPTION, + INTERFACE_DATA_DUPLEX, + INTERFACE_DATA_HELLO_TIME, + INTERFACE_DATA_LINK_UP, + INTERFACE_DATA_MAC, + INTERFACE_DATA_MAX_AGE, + INTERFACE_DATA_MULTICAST, + INTERFACE_DATA_PRIORITY, + INTERFACE_DATA_PROMISCUOUS, + INTERFACE_DATA_SPEED, + INTERFACE_DATA_STP, + INTERFACE_DATA_UP, + TRAFFIC_DATA_INTERFACE_ITEMS, + TRUE_STR, + WS_INTERFACES_KEY, +) +from ..common.enums import DeviceTypes, InterfaceTypes +from ..models.config_data import ConfigData +from ..models.edge_os_interface_data import EdgeOSInterfaceData +from .base_processor import BaseProcessor + +_LOGGER = logging.getLogger(__name__) + + +class InterfaceProcessor(BaseProcessor): + _interfaces: dict[str, EdgeOSInterfaceData] + + def __init__(self, config_data: ConfigData): + super().__init__(config_data) + + self.processor_type = DeviceTypes.INTERFACE + + self._interfaces: dict[str, EdgeOSInterfaceData] = {} + + def get_interfaces(self) -> list[str]: + return list(self._interfaces.keys()) + + def get_all(self) -> list[dict]: + items = [self._interfaces[item_key].to_dict() for item_key in self._interfaces] + + return items + + def get_interface(self, identifiers: set[tuple[str, str]]) -> dict | None: + interface: dict | None = None + interface_identifier = list(identifiers)[0][1] + + for interface_name in self._interfaces: + unique_id = self._get_device_info_unique_id(interface_name) + + if unique_id == interface_identifier: + interface = self._interfaces[interface_name].to_dict() + + return interface + + def get_data(self, interface_name: str) -> EdgeOSInterfaceData: + interface_data = self._interfaces.get(interface_name) + + return interface_data + + def get_device_info(self, item_id: str | None = None) -> DeviceInfo: + interface_name = item_id.upper() + + device_name = self._get_device_info_name(interface_name) + + unique_id = self._get_device_info_unique_id(interface_name) + + device_info = DeviceInfo( + identifiers={(DEFAULT_NAME, unique_id)}, + name=device_name, + model=self.processor_type, + manufacturer=DEFAULT_NAME, + via_device=(DEFAULT_NAME, self._hostname), + ) + + return device_info + + def _process_api_data(self): + super()._process_api_data() + + try: + system_section = self._api_data.get(API_DATA_SYSTEM, {}) + interface_types = system_section.get(API_DATA_INTERFACES, {}) + + for interface_type in interface_types: + interfaces = interface_types.get(interface_type) + + if interfaces is not None: + for interface_name in interfaces: + interface_data = interfaces.get(interface_name, {}) + int_type = InterfaceTypes(interface_type) + + self._extract_interface( + interface_name, interface_data, int_type + ) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract Interfaces data, Error: {ex}, Line: {line_number}" + ) + + def _process_ws_data(self): + try: + interfaces_data = self._ws_data.get(WS_INTERFACES_KEY, {}) + + for name in interfaces_data: + interface_item = self._interfaces.get(name) + stats = interfaces_data.get(name) + + if interface_item is None: + interface_data = interfaces_data.get(name) + interface_item = self._extract_interface( + name, interface_data, InterfaceTypes.DYNAMIC + ) + + if interface_item is not None: + self._update_interface_stats(interface_item, stats) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract WS data, Error: {ex}, Line: {line_number}" + ) + + def _extract_interface( + self, name: str, data: dict, interface_type: InterfaceTypes + ) -> EdgeOSInterfaceData: + interface = self._interfaces.get(name) + + try: + if data is not None: + if interface is None: + interface = EdgeOSInterfaceData(name, interface_type) + + interface.description = data.get(INTERFACE_DATA_DESCRIPTION) + interface.duplex = data.get(INTERFACE_DATA_DUPLEX) + interface.speed = data.get(INTERFACE_DATA_SPEED) + interface.bridge_group = data.get(INTERFACE_DATA_BRIDGE_GROUP) + interface.address = data.get(INTERFACE_DATA_ADDRESS) + interface.aging = data.get(INTERFACE_DATA_AGING) + interface.bridged_conntrack = data.get(INTERFACE_DATA_BRIDGED_CONNTRACK) + interface.hello_time = data.get(INTERFACE_DATA_HELLO_TIME) + interface.max_age = data.get(INTERFACE_DATA_MAX_AGE) + interface.priority = data.get(INTERFACE_DATA_PRIORITY) + interface.promiscuous = data.get(INTERFACE_DATA_PROMISCUOUS) + interface.stp = ( + data.get(INTERFACE_DATA_STP, FALSE_STR).lower() == TRUE_STR + ) + + self._interfaces[interface.unique_id] = interface + + if not interface.is_supported: + self._unique_log( + logging.INFO, + f"Ignoring interface {interface.name}, Interface type: {interface_type}", + ) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract interface data for {name}/{interface_type}, " + f"Error: {ex}, " + f"Line: {line_number}" + ) + + return interface + + @staticmethod + def _update_interface_stats(interface: EdgeOSInterfaceData, stats: dict): + try: + if stats is not None: + interface.up = ( + str(stats.get(INTERFACE_DATA_UP, False)).lower() == TRUE_STR + ) + interface.l1up = ( + str(stats.get(INTERFACE_DATA_LINK_UP, False)).lower() == TRUE_STR + ) + interface.mac = stats.get(INTERFACE_DATA_MAC) + interface.multicast = stats.get(INTERFACE_DATA_MULTICAST, 0) + interface.address = stats.get(ADDRESS_LIST, []) + + directions = [interface.received, interface.sent] + + for direction in directions: + stat_data = {} + for stat_key in TRAFFIC_DATA_INTERFACE_ITEMS: + key = f"{direction.direction}_{stat_key}" + stat_data_item = TRAFFIC_DATA_INTERFACE_ITEMS.get(stat_key) + + stat_data[stat_data_item] = float(stats.get(key)) + + direction.update(stat_data) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to update interface statistics for {interface.name}, " + f"Error: {ex}, " + f"Line: {line_number}" + ) diff --git a/custom_components/edgeos/data_processors/system_processor.py b/custom_components/edgeos/data_processors/system_processor.py new file mode 100644 index 0000000..171ddf3 --- /dev/null +++ b/custom_components/edgeos/data_processors/system_processor.py @@ -0,0 +1,229 @@ +from datetime import datetime +import logging +import sys + +from homeassistant.helpers.device_registry import DeviceInfo + +from ..common.consts import ( + API_DATA_DHCP_STATS, + API_DATA_SYS_INFO, + API_DATA_SYSTEM, + DATA_SYSTEM_SYSTEM, + DEFAULT_NAME, + DHCP_SERVER_LEASED, + DHCP_SERVER_STATS, + DISCOVER_DATA_FW_VERSION, + DISCOVER_DATA_PRODUCT, + FW_LATEST_STATE_CAN_UPGRADE, + SYSTEM_DATA_HOSTNAME, + SYSTEM_DATA_LOGIN, + SYSTEM_DATA_LOGIN_USER, + SYSTEM_DATA_LOGIN_USER_LEVEL, + SYSTEM_DATA_NTP, + SYSTEM_DATA_NTP_SERVER, + SYSTEM_DATA_OFFLOAD, + SYSTEM_DATA_OFFLOAD_HW_NAT, + SYSTEM_DATA_OFFLOAD_IPSEC, + SYSTEM_DATA_TIME_ZONE, + SYSTEM_DATA_TRAFFIC_ANALYSIS, + SYSTEM_DATA_TRAFFIC_ANALYSIS_DPI, + SYSTEM_DATA_TRAFFIC_ANALYSIS_EXPORT, + SYSTEM_INFO_DATA_FW_LATEST, + SYSTEM_INFO_DATA_FW_LATEST_STATE, + SYSTEM_INFO_DATA_FW_LATEST_URL, + SYSTEM_INFO_DATA_FW_LATEST_VERSION, + SYSTEM_INFO_DATA_SW_VER, + SYSTEM_STATS_DATA_CPU, + SYSTEM_STATS_DATA_MEM, + SYSTEM_STATS_DATA_UPTIME, + WS_DISCOVER_KEY, + WS_SYSTEM_STATS_KEY, +) +from ..common.enums import DeviceTypes +from ..models.config_data import ConfigData +from ..models.edge_os_system_data import EdgeOSSystemData +from .base_processor import BaseProcessor + +_LOGGER = logging.getLogger(__name__) + + +class SystemProcessor(BaseProcessor): + _system: EdgeOSSystemData | None = None + + def __init__(self, config_data: ConfigData): + super().__init__(config_data) + + self.processor_type = DeviceTypes.SYSTEM + + self._system = None + + def get(self) -> EdgeOSSystemData: + return self._system + + def get_device_info(self, item_id: str | None = None) -> DeviceInfo: + name = self._system.hostname.upper() + + device_info = DeviceInfo( + identifiers={(DEFAULT_NAME, name)}, + name=name, + model=self._system.product, + manufacturer=DEFAULT_NAME, + hw_version=self._system.fw_version, + ) + + return device_info + + def _process_api_data(self): + super()._process_api_data() + + try: + system_section = self._api_data.get(API_DATA_SYSTEM, {}) + system_info_section = self._api_data.get(API_DATA_SYS_INFO, {}) + + system_details = system_section.get(DATA_SYSTEM_SYSTEM, {}) + + system_data = EdgeOSSystemData() if self._system is None else self._system + + system_data.hostname = system_details.get(SYSTEM_DATA_HOSTNAME) + system_data.timezone = system_details.get(SYSTEM_DATA_TIME_ZONE) + + ntp: dict = system_details.get(SYSTEM_DATA_NTP, {}) + system_data.ntp_servers = ntp.get(SYSTEM_DATA_NTP_SERVER) + + offload: dict = system_details.get(SYSTEM_DATA_OFFLOAD, {}) + hardware_offload = EdgeOSSystemData.is_enabled( + offload, SYSTEM_DATA_OFFLOAD_HW_NAT + ) + ipsec_offload = EdgeOSSystemData.is_enabled( + offload, SYSTEM_DATA_OFFLOAD_IPSEC + ) + + system_data.hardware_offload = hardware_offload + system_data.ipsec_offload = ipsec_offload + + traffic_analysis: dict = system_details.get( + SYSTEM_DATA_TRAFFIC_ANALYSIS, {} + ) + dpi = EdgeOSSystemData.is_enabled( + traffic_analysis, SYSTEM_DATA_TRAFFIC_ANALYSIS_DPI + ) + traffic_analysis_export = EdgeOSSystemData.is_enabled( + traffic_analysis, SYSTEM_DATA_TRAFFIC_ANALYSIS_EXPORT + ) + + system_data.deep_packet_inspection = dpi + system_data.traffic_analysis_export = traffic_analysis_export + + sw_latest = system_info_section.get(SYSTEM_INFO_DATA_SW_VER) + fw_latest = system_info_section.get(SYSTEM_INFO_DATA_FW_LATEST, {}) + + fw_latest_state = fw_latest.get(SYSTEM_INFO_DATA_FW_LATEST_STATE) + fw_latest_version = fw_latest.get(SYSTEM_INFO_DATA_FW_LATEST_VERSION) + fw_latest_url = fw_latest.get(SYSTEM_INFO_DATA_FW_LATEST_URL) + + system_data.upgrade_available = ( + fw_latest_state == FW_LATEST_STATE_CAN_UPGRADE + ) + system_data.upgrade_url = fw_latest_url + system_data.upgrade_version = fw_latest_version + + system_data.sw_version = sw_latest + + login_details = system_details.get(SYSTEM_DATA_LOGIN, {}) + users = login_details.get(SYSTEM_DATA_LOGIN_USER, {}) + current_user = users.get(self._config_data.username, {}) + system_data.user_level = current_user.get(SYSTEM_DATA_LOGIN_USER_LEVEL) + + self._system = system_data + + message = ( + f"User {self._config_data.username} level is {self._system.user_level}, " + f"Interface status switch will not be created as it requires admin role" + ) + + self._unique_log(logging.INFO, message) + + self._update_leased_devices() + + warning_messages = [] + + if not self._system.deep_packet_inspection: + warning_messages.append("DPI (deep packet inspection) is turned off") + + if not self._system.traffic_analysis_export: + warning_messages.append("Traffic Analysis Export is turned off") + + if len(warning_messages) > 0: + warning_message = " and ".join(warning_messages) + + self._unique_log( + logging.WARNING, + f"Integration will not work correctly since {warning_message}", + ) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract System data, Error: {ex}, Line: {line_number}" + ) + + def _process_ws_data(self): + try: + system_stats_data = self._ws_data.get(WS_SYSTEM_STATS_KEY, {}) + discovery_data = self._ws_data.get(WS_DISCOVER_KEY, {}) + + system_data = self._system + + system_data.fw_version = discovery_data.get(DISCOVER_DATA_FW_VERSION) + system_data.product = discovery_data.get(DISCOVER_DATA_PRODUCT) + + uptime = float(system_stats_data.get(SYSTEM_STATS_DATA_UPTIME, 0)) + + system_data.cpu = int(system_stats_data.get(SYSTEM_STATS_DATA_CPU, 0)) + system_data.mem = int(system_stats_data.get(SYSTEM_STATS_DATA_MEM, 0)) + + if uptime != system_data.uptime: + system_data.uptime = uptime + system_data.last_reset = self._get_last_reset(uptime) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to update system statistics, " + f"Error: {ex}, " + f"Line: {line_number}" + ) + + def _update_leased_devices(self): + try: + unknown_devices = 0 + data_leases_stats_section = self._api_data.get(API_DATA_DHCP_STATS, {}) + + subnets = data_leases_stats_section.get(DHCP_SERVER_STATS, {}) + + for subnet in subnets: + subnet_data = subnets.get(subnet, {}) + unknown_devices += int(subnet_data.get(DHCP_SERVER_LEASED, 0)) + + self._system.leased_devices = unknown_devices + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract Unknown Devices data, Error: {ex}, Line: {line_number}" + ) + + @staticmethod + def _get_last_reset(uptime): + now = datetime.now().timestamp() + last_reset = now - uptime + + result = datetime.fromtimestamp(last_reset) + + return result diff --git a/custom_components/edgeos/device_tracker.py b/custom_components/edgeos/device_tracker.py index f2e2408..89515d5 100644 --- a/custom_components/edgeos/device_tracker.py +++ b/custom_components/edgeos/device_tracker.py @@ -1,33 +1,103 @@ """ Support for Ubiquiti EdgeOS routers. -HEAVILY based on the AsusWRT component For more details about this platform, please refer to the documentation at https://home-assistant.io/components/device_tracker.edgeos/ """ import logging -from .core.components.device_tracker import CoreScanner -from .core.helpers.setup_base_entry import async_setup_base_entry +from homeassistant.components.device_tracker import ( + ATTR_IP, + ATTR_MAC, + ScannerEntity, + SourceType, +) +from homeassistant.const import ATTR_ICON, Platform +from homeassistant.core import HomeAssistant + +from .common.base_entity import IntegrationBaseEntity, async_setup_base_entry +from .common.consts import ATTR_ATTRIBUTES, ATTR_HOSTNAME, ATTR_IS_ON +from .common.entity_descriptions import IntegrationDeviceTrackerEntityDescription +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator _LOGGER = logging.getLogger(__name__) -async def async_setup_entry(hass, config_entry, async_add_devices): +async def async_setup_entry(hass, entry, async_add_entities): """Set up the entity.""" await async_setup_base_entry( hass, - config_entry, - async_add_devices, - CoreScanner.get_domain(), - CoreScanner.get_component, + entry, + Platform.DEVICE_TRACKER, + IntegrationCoreScannerEntity, + async_add_entities, ) -async def async_unload_entry(hass, config_entry): - _LOGGER.info(f"Unload entry for {CoreScanner.get_domain()} domain: {config_entry}") +class IntegrationCoreScannerEntity(IntegrationBaseEntity, ScannerEntity): + """Represent a tracked device.""" + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationDeviceTrackerEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(hass, entity_description, coordinator, device_type, item_id) + + self._attr_ip_address: str | None = None + self._attr_mac_address: str | None = None + self._attr_source_type: SourceType | str | None = SourceType.ROUTER + self._attr_is_connected: bool = False + self._attr_hostname: str | None = None + + @property + def ip_address(self) -> str | None: + """Return the primary ip address of the device.""" + return self._attr_ip_address + + @property + def hostname(self) -> str | None: + """Return the hostname of the device.""" + return self._attr_hostname + + @property + def mac_address(self) -> str | None: + """Return the mac address of the device.""" + return self._attr_mac_address + + @property + def is_connected(self) -> bool: + """Return true if the device is connected to the network.""" + return self._attr_is_connected + + @property + def source_type(self) -> SourceType | str: + """Return the source type.""" + return self._attr_source_type + + def update_component(self, data): + """Fetch new state parameters for the sensor.""" + if data is not None: + is_connected = data.get(ATTR_IS_ON) + attributes = data.get(ATTR_ATTRIBUTES) + icon = data.get(ATTR_ICON) + + self._attr_is_connected = is_connected + self._attr_ip_address = attributes.get(ATTR_IP) + self._attr_mac_address = attributes.get(ATTR_MAC) + self._attr_hostname = attributes.get(ATTR_HOSTNAME) - return True + self._attr_extra_state_attributes = { + attribute: attributes[attribute] + for attribute in attributes + if attribute not in [ATTR_IP, ATTR_MAC] + } + if icon is not None: + self._attr_icon = icon -async def async_remove_entry(hass, entry) -> None: - _LOGGER.info(f"Remove entry for {CoreScanner.get_domain()} entry: {entry}") + else: + self._attr_is_connected = False diff --git a/custom_components/edgeos/diagnostics.py b/custom_components/edgeos/diagnostics.py index 96c6cde..49ed3c5 100644 --- a/custom_components/edgeos/diagnostics.py +++ b/custom_components/edgeos/diagnostics.py @@ -5,20 +5,13 @@ from typing import Any from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_PASSWORD from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.device_registry import DeviceEntry -from .component.helpers import get_ha -from .component.helpers.const import ( - API_DATA_INTERFACES, - API_DATA_SYSTEM, - DEVICE_LIST, - MESSAGES_COUNTER_SECTION, -) -from .component.managers.home_assistant import EdgeOSHomeAssistantManager -from .configuration.helpers.const import DOMAIN +from .common.consts import DEVICE_DATA_MAC, DOMAIN, INTERFACE_DATA_NAME +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator _LOGGER = logging.getLogger(__name__) @@ -29,107 +22,90 @@ async def async_get_config_entry_diagnostics( """Return diagnostics for a config entry.""" _LOGGER.debug("Starting diagnostic tool") - manager = get_ha(hass, entry.entry_id) + coordinator = hass.data[DOMAIN][entry.entry_id] - return _async_get_diagnostics(hass, manager, entry) + return _async_get_diagnostics(hass, coordinator, entry) async def async_get_device_diagnostics( hass: HomeAssistant, entry: ConfigEntry, device: DeviceEntry ) -> dict[str, Any]: """Return diagnostics for a device entry.""" - manager = get_ha(hass, entry.entry_id) + coordinator = hass.data[DOMAIN][entry.entry_id] - return _async_get_diagnostics(hass, manager, entry, device) + return _async_get_diagnostics(hass, coordinator, entry, device) @callback def _async_get_diagnostics( hass: HomeAssistant, - manager: EdgeOSHomeAssistantManager, + coordinator: Coordinator, entry: ConfigEntry, device: DeviceEntry | None = None, ) -> dict[str, Any]: """Return diagnostics for a config entry.""" _LOGGER.debug("Getting diagnostic information") - data = manager.config_data.to_dict() - configuration = manager.config_data.to_dict() - additional_data = manager.get_debug_data() + debug_data = coordinator.get_debug_data() - system = additional_data[API_DATA_SYSTEM] - device_list = additional_data[DEVICE_LIST] - interfaces = additional_data[API_DATA_INTERFACES] - data[MESSAGES_COUNTER_SECTION] = additional_data[MESSAGES_COUNTER_SECTION] - - if CONF_PASSWORD in configuration: - configuration.pop(CONF_PASSWORD) - - for configuration_key in configuration: - data[configuration_key] = configuration[configuration_key] - - data["disabled_by"] = entry.disabled_by - data["disabled_polling"] = entry.pref_disable_polling - - if CONF_PASSWORD in data: - data.pop(CONF_PASSWORD) + data = { + "disabled_by": entry.disabled_by, + "disabled_polling": entry.pref_disable_polling, + } if device: - device_name = next(iter(device.identifiers))[1] - - if device_name == manager.system_name: - data |= _async_device_as_dict(hass, system.to_dict(), manager.system_name) - - elif " Device " in device_name: - for network_device_id in device_list: - network_device = device_list.get(network_device_id) - - if manager.get_device_name(network_device) == device_name: - _LOGGER.debug( - f"Getting diagnostic information for device #{network_device.unique_id}" - ) - - data |= _async_device_as_dict( - hass, network_device.to_dict(), network_device.unique_id - ) + data["config"] = debug_data["config"] + data["data"] = debug_data["data"] + data["processors"] = debug_data["processors"] - break + device_data = coordinator.get_device_data(device.model, device.identifiers) - elif " Interface " in device_name: - for unique_id in interfaces: - interface = interfaces.get(unique_id) - - if manager.get_interface_name(interface) == device_name: - _LOGGER.debug( - f"Getting diagnostic information for interface #{interface.unique_id}" - ) - - data |= _async_device_as_dict( - hass, interface.to_dict(), interface.unique_id - ) + data |= _async_device_as_dict( + hass, + device.identifiers, + device_data, + ) - break else: _LOGGER.debug("Getting diagnostic information for all devices") + data = { + "config": debug_data["config"], + "data": debug_data["data"], + "processors": debug_data["processors"], + } + + processor_data = debug_data["processors"] + system_data = processor_data[DeviceTypes.SYSTEM] + device_data = processor_data[DeviceTypes.DEVICE] + interface_data = processor_data[DeviceTypes.INTERFACE] + data.update( devices=[ _async_device_as_dict( hass, - device_list[device_id].to_dict(), - device_list[device_id].unique_id, + coordinator.get_device_identifiers( + DeviceTypes.DEVICE, item.get(DEVICE_DATA_MAC) + ), + item, ) - for device_id in device_list + for item in device_data ], interfaces=[ _async_device_as_dict( hass, - interfaces[interface_id].to_dict(), - interfaces[interface_id].unique_id, + coordinator.get_device_identifiers( + DeviceTypes.INTERFACE, item.get(INTERFACE_DATA_NAME) + ), + item, ) - for interface_id in interfaces + for item in interface_data ], - system=_async_device_as_dict(hass, system.to_dict(), manager.system_name), + system=_async_device_as_dict( + hass, + coordinator.get_device_identifiers(DeviceTypes.SYSTEM), + system_data, + ), ) return data @@ -137,19 +113,22 @@ def _async_get_diagnostics( @callback def _async_device_as_dict( - hass: HomeAssistant, data: dict, unique_id: str + hass: HomeAssistant, identifiers, additional_data: dict ) -> dict[str, Any]: - """Represent a Shinobi monitor as a dictionary.""" + """Represent an EdgeOS based device as a dictionary.""" device_registry = dr.async_get(hass) entity_registry = er.async_get(hass) - ha_device = device_registry.async_get_device(identifiers={(DOMAIN, unique_id)}) + + ha_device = device_registry.async_get_device(identifiers=identifiers) + data = {} if ha_device: - data["home_assistant"] = { + data["device"] = { "name": ha_device.name, "name_by_user": ha_device.name_by_user, "disabled": ha_device.disabled, "disabled_by": ha_device.disabled_by, + "data": additional_data, "entities": [], } @@ -168,7 +147,7 @@ def _async_device_as_dict( # The context doesn't provide useful information in this case. state_dict.pop("context", None) - data["home_assistant"]["entities"].append( + data["device"]["entities"].append( { "disabled": entity_entry.disabled, "disabled_by": entity_entry.disabled_by, diff --git a/custom_components/edgeos/managers/config_manager.py b/custom_components/edgeos/managers/config_manager.py new file mode 100644 index 0000000..88f005c --- /dev/null +++ b/custom_components/edgeos/managers/config_manager.py @@ -0,0 +1,386 @@ +import json +import logging +import sys + +from cryptography.fernet import InvalidToken + +from homeassistant.config_entries import STORAGE_VERSION, ConfigEntry +from homeassistant.const import CONF_NAME, CONF_PASSWORD, CONF_USERNAME, Platform +from homeassistant.core import HomeAssistant +from homeassistant.helpers import translation +from homeassistant.helpers.device_registry import DeviceInfo +from homeassistant.helpers.json import JSONEncoder +from homeassistant.helpers.storage import Store + +from ..common.consts import ( + CONFIGURATION_FILE, + DEFAULT_CONSIDER_AWAY_INTERVAL, + DEFAULT_NAME, + DEFAULT_UPDATE_API_INTERVAL, + DEFAULT_UPDATE_ENTITIES_INTERVAL, + DOMAIN, + INVALID_TOKEN_SECTION, + STORAGE_DATA_CONSIDER_AWAY_INTERVAL, + STORAGE_DATA_LOG_INCOMING_MESSAGES, + STORAGE_DATA_MONITORED_DEVICES, + STORAGE_DATA_MONITORED_INTERFACES, + STORAGE_DATA_UPDATE_API_INTERVAL, + STORAGE_DATA_UPDATE_ENTITIES_INTERVAL, +) +from ..common.entity_descriptions import IntegrationEntityDescription +from ..models.config_data import ConfigData + +_LOGGER = logging.getLogger(__name__) + + +class ConfigManager: + _data: dict | None + _config_data: ConfigData + + _store: Store | None + _translations: dict | None + _password: str | None + _entry_title: str + _entry_id: str + + _is_set_up_mode: bool + _is_initialized: bool + + def __init__(self, hass: HomeAssistant | None, entry: ConfigEntry | None = None): + self._hass = hass + self._entry = entry + self._entry_id = None if entry is None else entry.entry_id + self._entry_title = DEFAULT_NAME if entry is None else entry.title + + self._config_data = ConfigData() + + self._data = None + + self._store = None + self._translations = None + + self._is_set_up_mode = entry is None + self._is_initialized = False + + if hass is not None: + self._store = Store( + hass, STORAGE_VERSION, CONFIGURATION_FILE, encoder=JSONEncoder + ) + + @property + def is_initialized(self) -> bool: + is_initialized = self._is_initialized + + return is_initialized + + @property + def entry_id(self) -> str: + entry_id = self._entry_id + + return entry_id + + @property + def entry_title(self) -> str: + entry_title = self._entry_title + + return entry_title + + @property + def entry(self) -> ConfigEntry: + entry = self._entry + + return entry + + @property + def monitored_interfaces(self): + result = self._data.get(STORAGE_DATA_MONITORED_INTERFACES, {}) + + return result + + @property + def monitored_devices(self): + result = self._data.get(STORAGE_DATA_MONITORED_DEVICES, {}) + + return result + + @property + def log_incoming_messages(self): + result = self._data.get(STORAGE_DATA_LOG_INCOMING_MESSAGES, False) + + return result + + @property + def consider_away_interval(self): + result = self._data.get( + STORAGE_DATA_CONSIDER_AWAY_INTERVAL, + DEFAULT_CONSIDER_AWAY_INTERVAL.total_seconds(), + ) + + return result + + @property + def update_entities_interval(self): + result = self._data.get( + STORAGE_DATA_UPDATE_ENTITIES_INTERVAL, + DEFAULT_UPDATE_ENTITIES_INTERVAL.total_seconds(), + ) + + return result + + @property + def update_api_interval(self): + result = self._data.get( + STORAGE_DATA_UPDATE_API_INTERVAL, + DEFAULT_UPDATE_API_INTERVAL.total_seconds(), + ) + + return result + + @property + def config_data(self) -> ConfigData: + config_data = self._config_data + + return config_data + + async def initialize(self, entry_config: dict): + try: + await self._load() + + self._config_data.update(entry_config) + + if self._hass is None: + self._translations = {} + + else: + self._translations = await translation.async_get_translations( + self._hass, self._hass.config.language, "entity", {DOMAIN} + ) + + _LOGGER.debug( + f"Translations loaded, Data: {json.dumps(self._translations)}" + ) + + self._is_initialized = True + + except InvalidToken: + self._is_initialized = False + + _LOGGER.error( + f"Invalid encryption key, Please follow instructions in {INVALID_TOKEN_SECTION}" + ) + + except Exception as ex: + self._is_initialized = False + + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to initialize configuration manager, Error: {ex}, Line: {line_number}" + ) + + def get_translation( + self, + platform: Platform, + entity_key: str, + attribute: str, + default_value: str | None = None, + ) -> str | None: + translation_key = ( + f"component.{DOMAIN}.entity.{platform}.{entity_key}.{attribute}" + ) + + translated_value = self._translations.get(translation_key, default_value) + + _LOGGER.debug( + "Translations requested, " + f"Key: {translation_key}, " + f"Default value: {default_value}, " + f"Value: {translated_value}" + ) + + return translated_value + + def get_entity_name( + self, + entity_description: IntegrationEntityDescription, + device_info: DeviceInfo, + ) -> str: + entity_key = entity_description.key + + device_name = device_info.get("name") + platform = entity_description.platform + + translated_name = self.get_translation( + platform, entity_key, CONF_NAME, entity_description.key + ) + + entity_name = ( + device_name + if translated_name is None or translated_name == "" + else f"{device_name} {translated_name}" + ) + + return entity_name + + def get_debug_data(self) -> dict: + data = self._config_data.to_dict() + + for key in self._data: + data[key] = self._data[key] + + return data + + def get_monitored_interface(self, interface_name: str): + is_enabled = self._data.get(STORAGE_DATA_MONITORED_INTERFACES, {}).get( + interface_name, False + ) + + return is_enabled + + def get_monitored_device(self, device_mac: str): + is_enabled = self._data.get(STORAGE_DATA_MONITORED_DEVICES, {}).get( + device_mac, False + ) + + return is_enabled + + async def _load(self): + self._data = None + + await self._load_config_from_file() + + _LOGGER.info(f"loaded: {self._data}") + should_save = False + + if self._data is None: + should_save = True + self._data = {} + + default_configuration = self._get_defaults() + _LOGGER.info(f"default_configuration: {default_configuration}") + + for key in default_configuration: + value = default_configuration[key] + + if key not in self._data: + _LOGGER.info(f"adding {key}") + should_save = True + self._data[key] = value + + if should_save: + _LOGGER.info("updated") + await self._save() + + @staticmethod + def _get_defaults() -> dict: + data = { + STORAGE_DATA_MONITORED_INTERFACES: {}, + STORAGE_DATA_MONITORED_DEVICES: {}, + STORAGE_DATA_LOG_INCOMING_MESSAGES: False, + STORAGE_DATA_CONSIDER_AWAY_INTERVAL: DEFAULT_CONSIDER_AWAY_INTERVAL.total_seconds(), + STORAGE_DATA_UPDATE_ENTITIES_INTERVAL: DEFAULT_UPDATE_ENTITIES_INTERVAL.total_seconds(), + STORAGE_DATA_UPDATE_API_INTERVAL: DEFAULT_UPDATE_API_INTERVAL.total_seconds(), + } + + return data + + async def _load_config_from_file(self): + if self._store is not None: + store_data = await self._store.async_load() + + if store_data is not None: + self._data = store_data.get(self._entry_id) + + async def remove(self, entry_id: str): + if self._store is None: + return + + store_data = await self._store.async_load() + + if store_data is not None and entry_id in store_data: + data = {key: store_data[key] for key in store_data} + data.pop(entry_id) + + await self._store.async_save(data) + + async def _save(self): + if self._store is None: + return + + should_save = False + store_data = await self._store.async_load() + + if store_data is None: + store_data = {} + + entry_data = store_data.get(self._entry_id, {}) + + _LOGGER.debug( + f"Storing config data: {json.dumps(self._data)}, " + f"Exiting: {json.dumps(entry_data)}" + ) + + for key in self._data: + stored_value = entry_data.get(key) + + if key in [CONF_PASSWORD, CONF_USERNAME]: + entry_data.pop(key) + + if stored_value is not None: + should_save = True + + else: + current_value = self._data.get(key) + + if stored_value != current_value: + should_save = True + + entry_data[key] = self._data[key] + + if should_save and self._entry_id is not None: + store_data[self._entry_id] = entry_data + + await self._store.async_save(store_data) + + async def set_monitored_interface(self, interface_name: str, is_enabled: bool): + _LOGGER.debug(f"Set monitored interface {interface_name} to {is_enabled}") + + self._data[STORAGE_DATA_MONITORED_INTERFACES][interface_name] = is_enabled + + await self._save() + + async def set_monitored_device(self, device_mac: str, is_enabled: bool): + _LOGGER.debug(f"Set monitored device {device_mac} to {is_enabled}") + + self._data[STORAGE_DATA_MONITORED_DEVICES][device_mac] = is_enabled + + await self._save() + + async def set_log_incoming_messages(self, enabled: bool): + _LOGGER.debug(f"Set log incoming messages to {enabled}") + + self._data[STORAGE_DATA_LOG_INCOMING_MESSAGES] = enabled + + await self._save() + + async def set_consider_away_interval(self, interval: int): + _LOGGER.debug(f"Changing {STORAGE_DATA_CONSIDER_AWAY_INTERVAL}: {interval}") + + self._data[STORAGE_DATA_CONSIDER_AWAY_INTERVAL] = interval + + await self._save() + + async def set_update_entities_interval(self, interval: int): + _LOGGER.debug(f"Changing {STORAGE_DATA_UPDATE_ENTITIES_INTERVAL}: {interval}") + + self._data[STORAGE_DATA_UPDATE_ENTITIES_INTERVAL] = interval + + await self._save() + + async def set_update_api_interval(self, interval: int): + _LOGGER.debug(f"Changing {STORAGE_DATA_UPDATE_API_INTERVAL}: {interval}") + + self._data[STORAGE_DATA_UPDATE_API_INTERVAL] = interval + + await self._save() diff --git a/custom_components/edgeos/managers/coordinator.py b/custom_components/edgeos/managers/coordinator.py new file mode 100644 index 0000000..338cdb6 --- /dev/null +++ b/custom_components/edgeos/managers/coordinator.py @@ -0,0 +1,879 @@ +from asyncio import sleep +from datetime import datetime, timedelta +import logging +import sys +from typing import Callable + +from homeassistant.components.device_tracker import ATTR_IP, ATTR_MAC +from homeassistant.components.homeassistant import SERVICE_RELOAD_CONFIG_ENTRY +from homeassistant.const import ATTR_STATE +from homeassistant.core import Event +from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.helpers.device_registry import DeviceInfo +from homeassistant.helpers.dispatcher import ( + async_dispatcher_connect, + async_dispatcher_send, +) +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +from ..common.connectivity_status import ConnectivityStatus +from ..common.consts import ( + ACTION_ENTITY_SET_NATIVE_VALUE, + ACTION_ENTITY_TURN_OFF, + ACTION_ENTITY_TURN_ON, + API_RECONNECT_INTERVAL, + ATTR_ACTIONS, + ATTR_ATTRIBUTES, + ATTR_HOSTNAME, + ATTR_IS_ON, + ATTR_LAST_ACTIVITY, + DOMAIN, + ENTITY_CONFIG_ENTRY_ID, + HA_NAME, + HEARTBEAT_INTERVAL, + SIGNAL_API_STATUS, + SIGNAL_DATA_CHANGED, + SIGNAL_DEVICE_ADDED, + SIGNAL_INTERFACE_ADDED, + SIGNAL_SYSTEM_ADDED, + SIGNAL_WS_STATUS, + SYSTEM_INFO_DATA_FW_LATEST_URL, + SYSTEM_INFO_DATA_FW_LATEST_VERSION, + WS_RECONNECT_INTERVAL, +) +from ..common.entity_descriptions import ( + ENTITY_DEVICE_MAPPING, + PLATFORMS, + IntegrationEntityDescription, +) +from ..common.enums import DeviceTypes, EntityKeys +from ..data_processors.base_processor import BaseProcessor +from ..data_processors.device_processor import DeviceProcessor +from ..data_processors.interface_processor import InterfaceProcessor +from ..data_processors.system_processor import SystemProcessor +from .config_manager import ConfigManager +from .rest_api import RestAPI +from .websockets import WebSockets + +_LOGGER = logging.getLogger(__name__) + + +class Coordinator(DataUpdateCoordinator): + """My custom coordinator.""" + + _api: RestAPI + _websockets: WebSockets | None + _processors: dict[DeviceTypes, BaseProcessor] | None = None + + _data_mapping: dict[ + str, + Callable[[IntegrationEntityDescription], dict | None] + | Callable[[IntegrationEntityDescription, str], dict | None], + ] | None + _system_status_details: dict | None + + _last_update: float + _last_heartbeat: float + + def __init__(self, hass, config_manager: ConfigManager): + """Initialize my coordinator.""" + super().__init__( + hass, + _LOGGER, + name=config_manager.entry_title, + update_interval=timedelta(seconds=config_manager.update_entities_interval), + update_method=self._async_update_data, + ) + + _LOGGER.debug("Initializing") + + entry = config_manager.entry + + signal_handlers = { + SIGNAL_API_STATUS: self._on_api_status_changed, + SIGNAL_WS_STATUS: self._on_ws_status_changed, + SIGNAL_DATA_CHANGED: self._on_data_changed, + } + + _LOGGER.debug(f"Registering signals for {signal_handlers.keys()}") + + for signal in signal_handlers: + handler = signal_handlers[signal] + + entry.async_on_unload(async_dispatcher_connect(hass, signal, handler)) + + config_data = config_manager.config_data + entry_id = config_manager.entry_id + + self._api = RestAPI(self.hass, config_data, entry_id) + + self._websockets = WebSockets(self.hass, config_data, entry_id) + + self._config_manager = config_manager + + self._data_mapping = None + + self._last_update = 0 + self._last_heartbeat = 0 + + self._can_load_components: bool = False + + self._system_processor = SystemProcessor(config_manager.config_data) + self._device_processor = DeviceProcessor(config_manager.config_data) + self._interface_processor = InterfaceProcessor(config_manager.config_data) + + self._discovered_objects = [] + + self._processors = { + DeviceTypes.SYSTEM: self._system_processor, + DeviceTypes.DEVICE: self._device_processor, + DeviceTypes.INTERFACE: self._interface_processor, + } + + _LOGGER.debug("Initializing done") + + @property + def api(self) -> RestAPI: + api = self._api + + return api + + @property + def websockets_data(self) -> dict: + data = self._websockets.data + + return data + + @property + def config_manager(self) -> ConfigManager: + config_manager = self._config_manager + + return config_manager + + async def on_home_assistant_start(self, _event_data: Event): + await self.initialize() + + async def initialize(self): + self._build_data_mapping() + + entry = self.config_manager.entry + await self.hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) + + _LOGGER.info(f"Start loading {DOMAIN} integration, Entry ID: {entry.entry_id}") + + await self.async_config_entry_first_refresh() + + await self._api.initialize() + + async def terminate(self): + await self._websockets.terminate() + + def get_debug_data(self) -> dict: + config_data = self._config_manager.get_debug_data() + + data = { + "config": config_data, + "data": { + "api": self._api.data, + "websockets": self._websockets.data, + }, + "processors": { + DeviceTypes.DEVICE: self._device_processor.get_all(), + DeviceTypes.INTERFACE: self._interface_processor.get_all(), + DeviceTypes.SYSTEM: self._system_processor.get().to_dict(), + }, + } + + return data + + async def _on_api_status_changed(self, entry_id: str, status: ConnectivityStatus): + if entry_id != self._config_manager.entry_id: + return + + if status == ConnectivityStatus.Connected: + await self._api.update() + + self._websockets.update_api_data( + self._api.data, self._config_manager.log_incoming_messages + ) + + await self._websockets.initialize() + + elif status in [ConnectivityStatus.Failed]: + await self._websockets.terminate() + + await sleep(API_RECONNECT_INTERVAL.total_seconds()) + + await self._api.initialize() + + elif status == ConnectivityStatus.InvalidCredentials: + self.update_interval = None + + async def _on_ws_status_changed(self, entry_id: str, status: ConnectivityStatus): + if entry_id != self._config_manager.entry_id: + return + + if status in [ConnectivityStatus.Failed, ConnectivityStatus.NotConnected]: + await self._websockets.terminate() + + await sleep(WS_RECONNECT_INTERVAL.total_seconds()) + + await self._api.initialize() + + def _on_system_discovered(self) -> None: + key = DeviceTypes.SYSTEM + + if key not in self._discovered_objects: + self._discovered_objects.append(key) + + async_dispatcher_send( + self.hass, + SIGNAL_SYSTEM_ADDED, + self._config_manager.entry_id, + DeviceTypes.SYSTEM, + ) + + def _on_device_discovered(self, device_mac: str) -> None: + key = f"{DeviceTypes.DEVICE} {device_mac}" + + if key not in self._discovered_objects: + self._discovered_objects.append(key) + + async_dispatcher_send( + self.hass, + SIGNAL_DEVICE_ADDED, + self._config_manager.entry_id, + DeviceTypes.DEVICE, + device_mac, + ) + + def _on_interface_discovered(self, interface_name: str) -> None: + key = f"{DeviceTypes.INTERFACE} {interface_name}" + + if key not in self._discovered_objects: + self._discovered_objects.append(key) + + async_dispatcher_send( + self.hass, + SIGNAL_INTERFACE_ADDED, + self._config_manager.entry_id, + DeviceTypes.INTERFACE, + interface_name, + ) + + async def _on_data_changed(self, entry_id: str): + if entry_id != self._config_manager.entry_id: + return + + api_connected = self._api.status == ConnectivityStatus.Connected + ws_client_connected = self._websockets.status == ConnectivityStatus.Connected + + is_ready = api_connected and ws_client_connected + + if is_ready: + for processor_type in self._processors: + processor = self._processors[processor_type] + processor.update(self._api.data, self._websockets.data) + + system = self._system_processor.get() + + if system.hostname is None: + return + + self._on_system_discovered() + + devices = self._device_processor.get_devices() + interfaces = self._interface_processor.get_interfaces() + + for interface_name in interfaces: + interface = self._interface_processor.get_data(interface_name) + + if interface.is_supported: + self._on_interface_discovered(interface_name) + + for device_mac in devices: + device = self._device_processor.get_data(device_mac) + + if not device.is_leased: + self._on_device_discovered(device_mac) + + async def _async_update_data(self): + """Fetch parameters from API endpoint. + + This is the place to pre-process the parameters to lookup tables + so entities can quickly look up their parameters. + """ + try: + _LOGGER.debug("Updating data") + + api_connected = self._api.status == ConnectivityStatus.Connected + ws_client_connected = ( + self._websockets.status == ConnectivityStatus.Connected + ) + + is_ready = api_connected and ws_client_connected + + if is_ready: + now = datetime.now().timestamp() + + if now - self._last_heartbeat >= HEARTBEAT_INTERVAL.total_seconds(): + await self._websockets.send_heartbeat() + + self._last_heartbeat = now + + if now - self._last_update >= self.config_manager.update_api_interval: + await self._api.update() + + self._last_update = now + + await self._on_data_changed(self.config_manager.entry_id) + + return {} + + except Exception as err: + raise UpdateFailed(f"Error communicating with API: {err}") + + def _build_data_mapping(self): + _LOGGER.debug("Building data mappers") + + data_mapping = { + EntityKeys.CPU_USAGE: self._get_cpu_usage_data, + EntityKeys.RAM_USAGE: self._get_ram_usage_data, + EntityKeys.FIRMWARE: self._get_firmware_data, + EntityKeys.LAST_RESTART: self._get_last_restart_data, + EntityKeys.UNKNOWN_DEVICES: self._get_unknown_devices_data, + EntityKeys.LOG_INCOMING_MESSAGES: self._get_log_incoming_messages_data, + EntityKeys.CONSIDER_AWAY_INTERVAL: self._get_consider_away_interval_data, + EntityKeys.UPDATE_ENTITIES_INTERVAL: self._get_update_entities_interval_data, + EntityKeys.UPDATE_API_INTERVAL: self._get_update_api_interval_data, + EntityKeys.INTERFACE_CONNECTED: self._get_interface_connected_data, + EntityKeys.INTERFACE_RECEIVED_DROPPED: self._get_interface_received_dropped_data, + EntityKeys.INTERFACE_SENT_DROPPED: self._get_interface_sent_dropped_data, + EntityKeys.INTERFACE_RECEIVED_ERRORS: self._get_interface_received_errors_data, + EntityKeys.INTERFACE_SENT_ERRORS: self._get_interface_sent_errors_data, + EntityKeys.INTERFACE_RECEIVED_PACKETS: self._get_interface_received_packets_data, + EntityKeys.INTERFACE_SENT_PACKETS: self._get_interface_sent_packets_data, + EntityKeys.INTERFACE_RECEIVED_RATE: self._get_interface_received_rate_data, + EntityKeys.INTERFACE_SENT_RATE: self._get_interface_sent_rate_data, + EntityKeys.INTERFACE_RECEIVED_TRAFFIC: self._get_interface_received_traffic_data, + EntityKeys.INTERFACE_SENT_TRAFFIC: self._get_interface_sent_traffic_data, + EntityKeys.INTERFACE_MONITORED: self._get_interface_monitored_data, + EntityKeys.INTERFACE_STATUS: self._get_interface_status_data, + EntityKeys.DEVICE_RECEIVED_RATE: self._get_device_received_rate_data, + EntityKeys.DEVICE_SENT_RATE: self._get_device_sent_rate_data, + EntityKeys.DEVICE_RECEIVED_TRAFFIC: self._get_device_received_traffic_data, + EntityKeys.DEVICE_SENT_TRAFFIC: self._get_device_sent_traffic_data, + EntityKeys.DEVICE_TRACKER: self._get_device_tracker_data, + EntityKeys.DEVICE_MONITORED: self._get_device_monitored_data, + } + + self._data_mapping = data_mapping + + _LOGGER.debug(f"Data retrieval mapping created, Mapping: {self._data_mapping}") + + def get_device_info( + self, + entity_description: IntegrationEntityDescription, + item_id: str | None = None, + ) -> DeviceInfo: + device_type = ENTITY_DEVICE_MAPPING.get(entity_description.key) + processor = self._processors[device_type] + + device_info = processor.get_device_info(item_id) + + return device_info + + def get_data( + self, + entity_description: IntegrationEntityDescription, + item_id: str | None = None, + ) -> dict | None: + result = None + + try: + handler = self._data_mapping.get(entity_description.key) + + if handler is None: + _LOGGER.warning( + f"Handler was not found for {entity_description.key}, Entity Description: {entity_description}" + ) + + else: + if item_id is None: + result = handler(entity_description) + + else: + result = handler(entity_description, item_id) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to extract data for {entity_description}, Error: {ex}, Line: {line_number}" + ) + + return result + + def get_device_identifiers( + self, device_type: DeviceTypes, item_id: str | None = None + ) -> set[tuple[str, str]]: + if device_type == DeviceTypes.DEVICE: + device_info = self._device_processor.get_device_info(item_id) + + elif device_type == DeviceTypes.INTERFACE: + device_info = self._interface_processor.get_device_info(item_id) + + else: + device_info = self._system_processor.get_device_info() + + identifiers = device_info.get("identifiers") + + return identifiers + + def get_device_data(self, model: str, identifiers: set[tuple[str, str]]): + if model == str(DeviceTypes.DEVICE): + device_data = self._device_processor.get_device(identifiers) + + elif model == str(DeviceTypes.INTERFACE): + device_data = self._interface_processor.get_interface(identifiers) + + else: + device_data = self._system_processor.get().to_dict() + + return device_data + + def get_device_action( + self, + entity_description: IntegrationEntityDescription, + monitor_id: str | None, + action_key: str, + ) -> Callable: + device_data = self.get_data(entity_description, monitor_id) + + actions = device_data.get(ATTR_ACTIONS) + async_action = actions.get(action_key) + + return async_action + + @staticmethod + def _get_date_time_from_timestamp(timestamp): + result = datetime.fromtimestamp(timestamp) + + return result + + def _get_cpu_usage_data(self, _entity_description) -> dict | None: + data = self._system_processor.get() + + result = { + ATTR_STATE: data.cpu, + } + + return result + + def _get_ram_usage_data(self, _entity_description) -> dict | None: + data = self._system_processor.get() + + result = { + ATTR_STATE: data.mem, + } + + return result + + def _get_firmware_data(self, _entity_description) -> dict | None: + data = self._system_processor.get() + + result = { + ATTR_IS_ON: data.upgrade_available, + ATTR_ATTRIBUTES: { + SYSTEM_INFO_DATA_FW_LATEST_URL: data.upgrade_url, + SYSTEM_INFO_DATA_FW_LATEST_VERSION: data.upgrade_version, + }, + } + + return result + + def _get_last_restart_data(self, _entity_description) -> dict | None: + data = self._system_processor.get() + + tz = datetime.now().astimezone().tzinfo + state = datetime.fromtimestamp(data.last_reset.timestamp(), tz=tz) + + result = {ATTR_STATE: state} + + return result + + def _get_unknown_devices_data(self, _entity_description) -> dict | None: + leased_devices = self._device_processor.get_leased_devices() + + result = { + ATTR_STATE: len(leased_devices.keys()), + ATTR_ATTRIBUTES: leased_devices, + } + + return result + + def _get_log_incoming_messages_data(self, _entity_description) -> dict | None: + result = { + ATTR_IS_ON: self.config_manager.log_incoming_messages, + ATTR_ACTIONS: { + ACTION_ENTITY_TURN_ON: self._set_log_incoming_messages_enabled, + ACTION_ENTITY_TURN_OFF: self._set_log_incoming_messages_disabled, + }, + } + + return result + + def _get_consider_away_interval_data(self, _entity_description) -> dict | None: + result = { + ATTR_STATE: self.config_manager.consider_away_interval, + ATTR_ACTIONS: { + ACTION_ENTITY_SET_NATIVE_VALUE: self._set_consider_away_interval, + }, + } + + return result + + def _get_update_entities_interval_data(self, _entity_description) -> dict | None: + result = { + ATTR_STATE: self.config_manager.update_entities_interval, + ATTR_ACTIONS: { + ACTION_ENTITY_SET_NATIVE_VALUE: self._set_update_entities_interval, + }, + } + + return result + + def _get_update_api_interval_data(self, _entity_description) -> dict | None: + result = { + ATTR_STATE: self.config_manager.update_api_interval, + ATTR_ACTIONS: { + ACTION_ENTITY_SET_NATIVE_VALUE: self._set_update_api_interval, + }, + } + + return result + + def _get_interface_connected_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_IS_ON: interface.l1up} + + return result + + def _get_interface_received_dropped_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.received.dropped} + + return result + + def _get_interface_sent_dropped_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.sent.dropped} + + return result + + def _get_interface_received_errors_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.received.errors} + + return result + + def _get_interface_sent_errors_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.sent.errors} + + return result + + def _get_interface_received_packets_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.received.packets} + + return result + + def _get_interface_sent_packets_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.sent.packets} + + return result + + def _get_interface_received_rate_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.received.total} + + return result + + def _get_interface_sent_rate_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.sent.rate} + + return result + + def _get_interface_received_traffic_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.received.total} + + return result + + def _get_interface_sent_traffic_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + + result = {ATTR_STATE: interface.sent.total} + + return result + + def _get_interface_monitored_data( + self, _entity_description, interface_name: str + ) -> dict | None: + state = self.config_manager.get_monitored_interface(interface_name) + + result = { + ATTR_IS_ON: state, + ATTR_ACTIONS: { + ACTION_ENTITY_TURN_ON: self._set_interface_monitor_enabled, + ACTION_ENTITY_TURN_OFF: self._set_interface_monitor_disabled, + }, + } + + return result + + def _get_interface_status_data( + self, _entity_description, interface_name: str + ) -> dict | None: + interface = self._interface_processor.get_data(interface_name) + interface_attributes = interface.get_attributes() + + result = { + ATTR_IS_ON: interface.up, + ATTR_ATTRIBUTES: interface_attributes, + ATTR_ACTIONS: { + ACTION_ENTITY_TURN_ON: self._set_interface_enabled, + ACTION_ENTITY_TURN_OFF: self._set_interface_disabled, + }, + } + + return result + + def _get_device_received_rate_data( + self, _entity_description, device_mac: str + ) -> dict | None: + device = self._device_processor.get_data(device_mac) + + result = {ATTR_STATE: device.received.rate} + + return result + + def _get_device_sent_rate_data( + self, _entity_description, device_mac: str + ) -> dict | None: + device = self._device_processor.get_data(device_mac) + + result = {ATTR_STATE: device.sent.rate} + + return result + + def _get_device_received_traffic_data( + self, _entity_description, device_mac: str + ) -> dict | None: + device = self._device_processor.get_data(device_mac) + + result = {ATTR_STATE: device.received.total} + + return result + + def _get_device_sent_traffic_data( + self, _entity_description, device_mac: str + ) -> dict | None: + device = self._device_processor.get_data(device_mac) + + result = {ATTR_STATE: device.sent.total} + + return result + + def _get_device_tracker_data( + self, _entity_description, device_mac: str + ) -> dict | None: + device = self._device_processor.get_data(device_mac) + consider_away_interval = self.config_manager.consider_away_interval + last_activity = self._get_date_time_from_timestamp(device.last_activity) + is_on = consider_away_interval >= device.last_activity_in_seconds + + result = { + ATTR_IS_ON: is_on, + ATTR_ATTRIBUTES: { + ATTR_LAST_ACTIVITY: last_activity, + ATTR_IP: device.ip, + ATTR_MAC: device.mac, + ATTR_HOSTNAME: device.hostname, + }, + } + + return result + + def _get_device_monitored_data( + self, _entity_description, device_mac: str + ) -> dict | None: + state = self.config_manager.get_monitored_device(device_mac) + device = self._device_processor.get_data(device_mac) + device_attributes = device.get_attributes() + + result = { + ATTR_IS_ON: state, + ATTR_ATTRIBUTES: device_attributes, + ATTR_ACTIONS: { + ACTION_ENTITY_TURN_ON: self._set_device_monitor_enabled, + ACTION_ENTITY_TURN_OFF: self._set_device_monitor_disabled, + }, + } + + return result + + async def _set_interface_enabled(self, _entity_description, interface_name: str): + _LOGGER.debug(f"Enable interface {interface_name}") + interface = self._interface_processor.get_data(interface_name) + + await self._api.set_interface_state(interface, True) + + async def _set_interface_disabled(self, _entity_description, interface_name: str): + _LOGGER.debug(f"Disable interface {interface_name}") + interface = self._interface_processor.get_data(interface_name) + + await self._api.set_interface_state(interface, False) + + async def _set_interface_monitor_enabled( + self, _entity_description, interface_name: str + ): + _LOGGER.debug(f"Enable monitoring for interface {interface_name}") + + await self._config_manager.set_monitored_interface(interface_name, True) + + self._remove_entities_of_device(DeviceTypes.INTERFACE, interface_name) + + async def _set_interface_monitor_disabled( + self, _entity_description, interface_name: str + ): + _LOGGER.debug(f"Disable monitoring for interface {interface_name}") + + await self._config_manager.set_monitored_interface(interface_name, False) + + self._remove_entities_of_device(DeviceTypes.INTERFACE, interface_name) + + async def _set_device_monitor_enabled(self, _entity_description, device_mac: str): + _LOGGER.debug(f"Enable monitoring for device {device_mac}") + + await self._config_manager.set_monitored_device(device_mac, True) + + self._remove_entities_of_device(DeviceTypes.DEVICE, device_mac) + + async def _set_device_monitor_disabled(self, _entity_description, device_mac: str): + _LOGGER.debug(f"Disable monitoring for device {device_mac}") + + await self._config_manager.set_monitored_device(device_mac, False) + + self._remove_entities_of_device(DeviceTypes.DEVICE, device_mac) + + def _remove_entities_of_device(self, device_type: DeviceTypes, item_id: str): + key = f"{device_type} {item_id}" + + if device_type == DeviceTypes.DEVICE: + device_info = self._device_processor.get_device_info(item_id) + + elif device_type == DeviceTypes.INTERFACE: + device_info = self._interface_processor.get_device_info(item_id) + + else: + return + + entity_registry = er.async_get(self.hass) + device_registry = dr.async_get(self.hass) + + device_info_identifier = device_info.get("identifiers") + device_data = device_registry.async_get_device( + identifiers=device_info_identifier + ) + device_id = device_data.id + + entities = entity_registry.entities.get_entries_for_device_id(device_id) + for entity in entities: + entity_registry.async_remove(entity.entity_id) + + self._discovered_objects.remove(key) + + if device_type == DeviceTypes.DEVICE: + self._on_device_discovered(item_id) + + elif device_type == DeviceTypes.INTERFACE: + self._on_interface_discovered(item_id) + + async def _set_log_incoming_messages_enabled(self, _entity_description): + _LOGGER.debug("Enable log incoming messages") + + await self._config_manager.set_log_incoming_messages(True) + + self._websockets.update_api_data( + self._api.data, self.config_manager.log_incoming_messages + ) + + async def _set_log_incoming_messages_disabled(self, _entity_description): + _LOGGER.debug("Disable log incoming messages") + + await self._config_manager.set_log_incoming_messages(False) + + self._websockets.update_api_data( + self._api.data, self.config_manager.log_incoming_messages + ) + + async def _set_consider_away_interval(self, _entity_description, value: int): + _LOGGER.debug("Disable log incoming messages") + + await self._config_manager.set_consider_away_interval(value) + + async def _set_update_entities_interval(self, _entity_description, value: int): + _LOGGER.debug("Disable log incoming messages") + + await self._config_manager.set_update_entities_interval(value) + + await self._reload_integration() + + async def _set_update_api_interval(self, _entity_description, value: int): + _LOGGER.debug("Disable log incoming messages") + + await self._config_manager.set_update_api_interval(value) + + await self._reload_integration() + + async def _reload_integration(self): + data = {ENTITY_CONFIG_ENTRY_ID: self.config_manager.entry_id} + + await self.hass.services.async_call(HA_NAME, SERVICE_RELOAD_CONFIG_ENTRY, data) diff --git a/custom_components/edgeos/managers/flow_manager.py b/custom_components/edgeos/managers/flow_manager.py new file mode 100644 index 0000000..65b2c4e --- /dev/null +++ b/custom_components/edgeos/managers/flow_manager.py @@ -0,0 +1,139 @@ +"""Config flow to configure.""" +from __future__ import annotations + +from copy import copy +import logging +from typing import Any + +from cryptography.fernet import InvalidToken + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.data_entry_flow import FlowHandler + +from ..common.connectivity_status import ConnectivityStatus +from ..common.consts import CONF_TITLE, DEFAULT_NAME +from ..models.config_data import DATA_KEYS, ConfigData +from ..models.exceptions import LoginError +from .config_manager import ConfigManager +from .password_manager import PasswordManager +from .rest_api import RestAPI + +_LOGGER = logging.getLogger(__name__) + + +class IntegrationFlowManager: + _hass: HomeAssistant + _entry: ConfigEntry | None + + _flow_handler: FlowHandler + _flow_id: str + + _config_manager: ConfigManager + + def __init__( + self, + hass: HomeAssistant, + flow_handler: FlowHandler, + entry: ConfigEntry | None = None, + ): + self._hass = hass + self._flow_handler = flow_handler + self._entry = entry + self._flow_id = "user" if entry is None else "init" + self._config_manager = ConfigManager(self._hass, None) + + async def async_step(self, user_input: dict | None = None): + """Manage the domain options.""" + _LOGGER.info(f"Config flow started, Step: {self._flow_id}, Input: {user_input}") + + form_errors = None + + if user_input is None: + if self._entry is None: + user_input = {} + + else: + user_input = {key: self._entry.data[key] for key in self._entry.data} + user_input[CONF_TITLE] = self._entry.title + + _LOGGER.info(user_input) + + await PasswordManager.decrypt( + self._hass, user_input, self._entry.entry_id + ) + + else: + try: + await self._config_manager.initialize(user_input) + config_data = ConfigData() + config_data.update(user_input) + + api = RestAPI(self._hass, config_data) + + await api.validate() + + if api.status == ConnectivityStatus.Connected: + _LOGGER.debug("User inputs are valid") + + if self._entry is None: + data = copy(user_input) + + else: + data = await self.remap_entry_data(user_input) + + await PasswordManager.encrypt(self._hass, data) + + title = data.get(CONF_TITLE, DEFAULT_NAME) + + if CONF_TITLE in data: + data.pop(CONF_TITLE) + + return self._flow_handler.async_create_entry(title=title, data=data) + + else: + error_key = ConnectivityStatus.get_ha_error(api.status) + + except LoginError: + error_key = "invalid_admin_credentials" + + except InvalidToken: + error_key = "corrupted_encryption_key" + + if error_key is not None: + form_errors = {"base": error_key} + + _LOGGER.warning(f"Failed to create integration, Error Key: {error_key}") + + schema = ConfigData.default_schema(user_input) + + return self._flow_handler.async_show_form( + step_id=self._flow_id, data_schema=schema, errors=form_errors + ) + + async def remap_entry_data(self, options: dict[str, Any]) -> dict[str, Any]: + config_options = {} + config_data = {} + + entry = self._entry + entry_data = entry.data + + title = DEFAULT_NAME + + for key in options: + if key in DATA_KEYS: + config_data[key] = options.get(key, entry_data.get(key)) + + elif key == CONF_TITLE: + title = options.get(key, DEFAULT_NAME) + + else: + config_options[key] = options.get(key) + + await PasswordManager.encrypt(self._hass, config_data) + + self._hass.config_entries.async_update_entry( + entry, data=config_data, title=title + ) + + return config_options diff --git a/custom_components/edgeos/managers/password_manager.py b/custom_components/edgeos/managers/password_manager.py new file mode 100644 index 0000000..498479e --- /dev/null +++ b/custom_components/edgeos/managers/password_manager.py @@ -0,0 +1,129 @@ +import logging +import sys + +from cryptography.fernet import Fernet, InvalidToken + +from homeassistant.config_entries import STORAGE_VERSION +from homeassistant.const import CONF_PASSWORD +from homeassistant.core import HomeAssistant +from homeassistant.helpers.json import JSONEncoder +from homeassistant.helpers.storage import Store + +from ..common.consts import CONFIGURATION_FILE, INVALID_TOKEN_SECTION, STORAGE_DATA_KEY + +_LOGGER = logging.getLogger(__name__) + + +class PasswordManager: + _encryption_key: str | None + _crypto: Fernet | None + _entry_id: str + + def __init__(self, hass: HomeAssistant | None, entry_id: str = ""): + self._hass = hass + self._entry_id = entry_id + + self._encryption_key = None + self._crypto = None + + if hass is None: + self._store = None + + else: + self._store = Store( + hass, STORAGE_VERSION, CONFIGURATION_FILE, encoder=JSONEncoder + ) + + async def initialize(self): + try: + await self._load_encryption_key() + + except InvalidToken: + _LOGGER.error( + f"Invalid encryption key, Please follow instructions in {INVALID_TOKEN_SECTION}" + ) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to initialize configuration manager, Error: {ex}, Line: {line_number}" + ) + + @staticmethod + async def decrypt(hass: HomeAssistant, data: dict, entry_id: str = "") -> None: + instance = PasswordManager(hass, entry_id) + + await instance.initialize() + + password = data.get(CONF_PASSWORD) + password_decrypted = instance._decrypt(password) + data[CONF_PASSWORD] = password_decrypted + + @staticmethod + async def encrypt(hass: HomeAssistant, data: dict, entry_id: str = "") -> None: + instance = PasswordManager(hass, entry_id) + + await instance.initialize() + + if CONF_PASSWORD in data: + password = data.get(CONF_PASSWORD) + password_encrypted = instance._encrypt(password) + + data[CONF_PASSWORD] = password_encrypted + + async def _load_encryption_key(self): + store_data = None + + if self._store is not None: + store_data = await self._store.async_load() + + if store_data is not None: + if STORAGE_DATA_KEY in store_data: + self._encryption_key = store_data.get(STORAGE_DATA_KEY) + + else: + for store_data_key in store_data: + if store_data_key == self._entry_id: + entry_configuration = store_data[store_data_key] + + if STORAGE_DATA_KEY in entry_configuration: + self._encryption_key = entry_configuration.get( + STORAGE_DATA_KEY + ) + + entry_configuration.pop(STORAGE_DATA_KEY) + + if self._encryption_key is None: + self._encryption_key = Fernet.generate_key().decode("utf-8") + + await self._save() + + self._crypto = Fernet(self._encryption_key.encode()) + + async def _save(self): + if self._store is None: + return + + store_data = await self._store.async_load() + + if store_data is None: + store_data = {} + + if store_data.get(STORAGE_DATA_KEY) != self._encryption_key: + store_data[STORAGE_DATA_KEY] = self._encryption_key + + await self._store.async_save(store_data) + + def _encrypt(self, data: str) -> str: + if data is not None: + data = self._crypto.encrypt(data.encode()).decode() + + return data + + def _decrypt(self, data: str) -> str: + if data is not None and len(data) > 0: + data = self._crypto.decrypt(data.encode()).decode() + + return data diff --git a/custom_components/edgeos/component/api/api.py b/custom_components/edgeos/managers/rest_api.py similarity index 71% rename from custom_components/edgeos/component/api/api.py rename to custom_components/edgeos/managers/rest_api.py index 6bf8dbb..cb56207 100644 --- a/custom_components/edgeos/component/api/api.py +++ b/custom_components/edgeos/managers/rest_api.py @@ -1,29 +1,21 @@ from __future__ import annotations from asyncio import sleep -from collections.abc import Awaitable, Callable from datetime import datetime, timedelta import json import logging import sys +from typing import Any -from aiohttp import CookieJar +from aiohttp import ClientSession, CookieJar from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.core import HomeAssistant +from homeassistant.helpers.aiohttp_client import async_create_clientsession +from homeassistant.helpers.dispatcher import async_dispatcher_send -from ...configuration.helpers.const import ( - COOKIE_BEAKER_SESSION_ID, - COOKIE_CSRF_TOKEN, - COOKIE_PHPSESSID, - HEADER_CSRF_TOKEN, - MAXIMUM_RECONNECT, -) -from ...configuration.models.config_data import ConfigData -from ...core.api.base_api import BaseAPI -from ...core.helpers.const import EMPTY_STRING -from ...core.helpers.enums import ConnectivityStatus -from ..helpers.const import ( +from ..common.connectivity_status import ConnectivityStatus +from ..common.consts import ( API_DATA, API_DATA_COOKIES, API_DATA_INTERFACES, @@ -42,49 +34,91 @@ API_URL_PARAMETER_BASE_URL, API_URL_PARAMETER_SUBSET, API_URL_PARAMETER_TIMESTAMP, + COOKIE_BEAKER_SESSION_ID, + COOKIE_CSRF_TOKEN, + COOKIE_PHPSESSID, + DEFAULT_NAME, + EMPTY_STRING, + HEADER_CSRF_TOKEN, HEARTBEAT_MAX_AGE, + MAXIMUM_RECONNECT, RESPONSE_ERROR_KEY, RESPONSE_FAILURE_CODE, RESPONSE_OUTPUT, RESPONSE_SUCCESS_KEY, + SIGNAL_API_STATUS, + SIGNAL_DATA_CHANGED, STRING_DASH, STRING_UNDERSCORE, SYSTEM_DATA_DISABLE, TRUE_STR, UPDATE_DATE_ENDPOINTS, ) +from ..models.config_data import ConfigData from ..models.edge_os_interface_data import EdgeOSInterfaceData from ..models.exceptions import SessionTerminatedException _LOGGER = logging.getLogger(__name__) -class IntegrationAPI(BaseAPI): - """The Class for handling the data retrieval.""" +class RestAPI: + _hass: HomeAssistant | None + data: dict - _config_data: ConfigData | None + _base_url: str | None + _status: ConnectivityStatus | None + _session: ClientSession | None + _entry_id: str | None + _dispatched_devices: list + _dispatched_server: bool + + _last_valid: datetime | None def __init__( - self, - hass: HomeAssistant | None, - async_on_data_changed: Callable[[], Awaitable[None]] | None = None, - async_on_status_changed: Callable[[ConnectivityStatus], Awaitable[None]] - | None = None, + self, hass: HomeAssistant, config_data: ConfigData, entry_id: str | None = None ): - super().__init__(hass, async_on_data_changed, async_on_status_changed) - try: - self._config_data = None - self._cookies = {} - self._last_valid = None + self._hass = hass + self._support_video_browser_api = False self.data = {} + self._cookies = {} + + self._config_data = config_data + + self._local_async_dispatcher_send = None + + self._status = None + + self._session = None + self._entry_id = entry_id + self._dispatched_devices = [] + self._dispatched_server = False + self._last_valid = None except Exception as ex: exc_type, exc_obj, tb = sys.exc_info() line_number = tb.tb_lineno - _LOGGER.error(f"Failed to load API, error: {ex}, line: {line_number}") + _LOGGER.error( + f"Failed to load {DEFAULT_NAME} API, error: {ex}, line: {line_number}" + ) + + @property + def is_connected(self): + result = self._session is not None + + return result + + @property + def status(self) -> str | None: + status = self._status + + return status + + @property + def _is_home_assistant(self): + return self._hass is not None @property def session_id(self): @@ -108,26 +142,24 @@ def csrf_token(self): def cookies_data(self): return self._cookies - async def initialize(self, config_data: ConfigData): - _LOGGER.info("Initializing API") + async def _do_nothing(self, _status: ConnectivityStatus): + pass - try: - self._config_data = config_data + async def initialize(self): + _LOGGER.info("Initializing EdgeOS API") - cookie_jar = CookieJar(unsafe=True) + self._set_status(ConnectivityStatus.Connecting) - await self.initialize_session(cookies=self._cookies, cookie_jar=cookie_jar) + cookie_jar = CookieJar(unsafe=True) - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno + await self._initialize_session(cookie_jar) - _LOGGER.error(f"Failed to initialize API, error: {ex}, line: {line_number}") + await self.login() - async def validate(self, data: dict | None = None): - config_data = ConfigData.from_dict(data) + async def validate(self): + await self.initialize() - await self.initialize(config_data) + await self.login() def _get_cookie_data(self, cookie_key): cookie_data = None @@ -137,69 +169,23 @@ def _get_cookie_data(self, cookie_key): return cookie_data - async def login(self): - await super().login() - - try: - username = self._config_data.username - password = self._config_data.password - - credentials = {CONF_USERNAME: username, CONF_PASSWORD: password} - - url = self._config_data.url - - if self.session.closed: - raise SessionTerminatedException() - - async with self.session.post(url, data=credentials, ssl=False) as response: - all_cookies = self.session.cookie_jar.filter_cookies(response.url) - - for key, cookie in all_cookies.items(): - self._cookies[cookie.key] = cookie.value - - response.raise_for_status() - - logged_in = ( - self.beaker_session_id is not None - and self.beaker_session_id == self.session_id - ) - - if logged_in: - html = await response.text() - html_lines = html.splitlines() - for line in html_lines: - if "EDGE.DeviceModel" in line: - line_parts = line.split(" = ") - value = line_parts[len(line_parts) - 1] - self.data[API_DATA_PRODUCT] = value.replace( - "'", EMPTY_STRING - ) - self.data[API_DATA_SESSION_ID] = self.session_id - self.data[API_DATA_COOKIES] = self._cookies - - await self.set_status(ConnectivityStatus.Connected) - - break - else: - _LOGGER.error("Failed to login, Invalid credentials") - - if self.beaker_session_id is None and self.session_id is not None: - await self.set_status(ConnectivityStatus.Failed) - else: - await self.set_status(ConnectivityStatus.InvalidCredentials) - - except SessionTerminatedException: - await self.set_status(ConnectivityStatus.Disconnected) - - raise SessionTerminatedException() - - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno + def _build_endpoint( + self, + endpoint, + timestamp: str | None = None, + action: str | None = None, + subset: str | None = None, + ): + data = { + API_URL_PARAMETER_BASE_URL: self._config_data.api_url, + API_URL_PARAMETER_TIMESTAMP: timestamp, + API_URL_PARAMETER_ACTION: action, + API_URL_PARAMETER_SUBSET: subset, + } - _LOGGER.error(f"Failed to login, Error: {ex}, Line: {line_number}") + url = endpoint.format(**data) - await self.set_status(ConnectivityStatus.NotFound) + return url async def _async_get( self, @@ -222,8 +208,8 @@ async def _async_get( retry_attempt = retry_attempt + 1 try: - if self.session is not None: - async with self.session.get(url, ssl=False) as response: + if self._session is not None: + async with self._session.get(url, ssl=False) as response: status = response.status message = ( @@ -234,7 +220,7 @@ async def _async_get( result = await response.json() break elif status == 403: - self.session = None + self._session = None self._cookies = {} break @@ -251,14 +237,14 @@ async def _async_get( _LOGGER.warning(f"Request failed, {message}") - await self.set_status(ConnectivityStatus.Disconnected) + self._set_status(ConnectivityStatus.Disconnected) return result def _get_post_headers(self): headers = {} - for header_key in self.session.headers: - header = self.session.headers.get(header_key) + for header_key in self._session.headers: + header = self._session.headers.get(header_key) if header is not None: headers[header_key] = header @@ -273,11 +259,11 @@ async def _async_post(self, endpoint, data): try: url = self._build_endpoint(endpoint) - if self.session is not None: + if self._session is not None: headers = self._get_post_headers() data_json = json.dumps(data) - async with self.session.post( + async with self._session.post( url, headers=headers, data=data_json, ssl=False ) as response: response.raise_for_status() @@ -293,6 +279,129 @@ async def _async_post(self, endpoint, data): return result + async def update(self): + _LOGGER.debug(f"Updating data from device ({self._config_data.hostname})") + + if self.status == ConnectivityStatus.Failed: + await self.initialize() + + if self.status == ConnectivityStatus.Connected: + await self._load_system_data() + + for endpoint in UPDATE_DATE_ENDPOINTS: + await self._load_general_data(endpoint) + + self.data[API_DATA_LAST_UPDATE] = datetime.now().isoformat() + + self._async_dispatcher_send(SIGNAL_DATA_CHANGED) + + async def login(self): + try: + username = self._config_data.username + password = self._config_data.password + + credentials = {CONF_USERNAME: username, CONF_PASSWORD: password} + + url = self._config_data.api_url + + if self._session.closed: + raise SessionTerminatedException() + + async with self._session.post(url, data=credentials, ssl=False) as response: + all_cookies = self._session.cookie_jar.filter_cookies(response.url) + + for key, cookie in all_cookies.items(): + self._cookies[cookie.key] = cookie.value + + response.raise_for_status() + + logged_in = ( + self.beaker_session_id is not None + and self.beaker_session_id == self.session_id + ) + + if logged_in: + html = await response.text() + html_lines = html.splitlines() + for line in html_lines: + if "EDGE.DeviceModel" in line: + line_parts = line.split(" = ") + value = line_parts[len(line_parts) - 1] + self.data[API_DATA_PRODUCT] = value.replace( + "'", EMPTY_STRING + ) + self.data[API_DATA_SESSION_ID] = self.session_id + self.data[API_DATA_COOKIES] = self._cookies + + self._set_status(ConnectivityStatus.Connected) + + break + else: + _LOGGER.error("Failed to login, Invalid credentials") + + if self.beaker_session_id is None and self.session_id is not None: + self._set_status(ConnectivityStatus.Failed) + else: + self._set_status(ConnectivityStatus.InvalidCredentials) + + except SessionTerminatedException: + self._set_status(ConnectivityStatus.Disconnected) + + raise SessionTerminatedException() + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error(f"Failed to login, Error: {ex}, Line: {line_number}") + + self._set_status(ConnectivityStatus.NotFound) + + async def _initialize_session(self, cookie_jar=None): + try: + if self._is_home_assistant: + self._session = async_create_clientsession( + hass=self._hass, cookies=self._cookies, cookie_jar=cookie_jar + ) + + else: + self._session = ClientSession( + cookies=self._cookies, cookie_jar=cookie_jar + ) + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.warning( + f"Failed to initialize session, Error: {str(ex)}, Line: {line_number}" + ) + + self._set_status(ConnectivityStatus.Failed) + + def _set_status(self, status: ConnectivityStatus): + if status != self._status: + log_level = ConnectivityStatus.get_log_level(status) + + _LOGGER.log( + log_level, + f"Status changed from '{self._status}' to '{status}'", + ) + + self._status = status + + self._async_dispatcher_send(SIGNAL_API_STATUS, status) + + def set_local_async_dispatcher_send(self, callback): + self._local_async_dispatcher_send = callback + + def _async_dispatcher_send(self, signal: str, *args: Any) -> None: + if self._hass is None: + self._local_async_dispatcher_send(signal, None, *args) + + else: + async_dispatcher_send(self._hass, signal, self._entry_id, *args) + async def async_send_heartbeat(self, max_age=HEARTBEAT_MAX_AGE): ts = None @@ -326,25 +435,7 @@ async def async_send_heartbeat(self, max_age=HEARTBEAT_MAX_AGE): is_valid = ts is not None and self._last_valid == ts if not is_valid: - await self.set_status(ConnectivityStatus.Disconnected) - - async def async_update(self): - try: - await self._load_system_data() - - for endpoint in UPDATE_DATE_ENDPOINTS: - await self._load_general_data(endpoint) - - self.data[API_DATA_LAST_UPDATE] = datetime.now().isoformat() - - await self.fire_data_changed_event() - except Exception as ex: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno - - _LOGGER.error( - f"Failed to extract WS data, Error: {ex}, Line: {line_number}" - ) + self._set_status(ConnectivityStatus.Disconnected) async def _load_system_data(self): try: @@ -439,21 +530,3 @@ async def set_interface_state( _LOGGER.error( f"Failed to set state of interface {interface.name} to {is_enabled}" ) - - def _build_endpoint( - self, - endpoint, - timestamp: str | None = None, - action: str | None = None, - subset: str | None = None, - ): - data = { - API_URL_PARAMETER_BASE_URL: self._config_data.url, - API_URL_PARAMETER_TIMESTAMP: timestamp, - API_URL_PARAMETER_ACTION: action, - API_URL_PARAMETER_SUBSET: subset, - } - - url = endpoint.format(**data) - - return url diff --git a/custom_components/edgeos/component/api/websocket.py b/custom_components/edgeos/managers/websockets.py similarity index 69% rename from custom_components/edgeos/component/api/websocket.py rename to custom_components/edgeos/managers/websockets.py index 5375038..bd5b24b 100644 --- a/custom_components/edgeos/component/api/websocket.py +++ b/custom_components/edgeos/managers/websockets.py @@ -1,36 +1,39 @@ -""" -websocket. -""" from __future__ import annotations -from collections.abc import Awaitable, Callable +import asyncio +from collections.abc import Awaitable +from datetime import datetime import json import logging import re import sys -from urllib.parse import urlparse +from typing import Any, Callable import aiohttp +from aiohttp import ClientSession from homeassistant.core import HomeAssistant +from homeassistant.helpers.aiohttp_client import async_create_clientsession +from homeassistant.helpers.dispatcher import async_dispatcher_send -from ...configuration.helpers.const import WEBSOCKET_URL_TEMPLATE -from ...configuration.models.config_data import ConfigData -from ...core.api.base_api import BaseAPI -from ...core.helpers.const import EMPTY_STRING -from ...core.helpers.enums import ConnectivityStatus -from ..helpers.const import ( +from ..common.connectivity_status import ConnectivityStatus +from ..common.consts import ( ADDRESS_HW_ADDR, ADDRESS_IPV4, ADDRESS_LIST, API_DATA_COOKIES, + API_DATA_LAST_UPDATE, API_DATA_SESSION_ID, BEGINS_WITH_SIX_DIGITS, DEVICE_LIST, + DISCONNECT_INTERVAL, DISCOVER_DEVICE_ITEMS, + EMPTY_STRING, INTERFACE_DATA_MULTICAST, INTERFACES_MAIN_MAP, INTERFACES_STATS, + SIGNAL_DATA_CHANGED, + SIGNAL_WS_STATUS, STRING_COLON, STRING_COMMA, TRAFFIC_DATA_DEVICE_ITEMS, @@ -51,33 +54,77 @@ WS_TOPIC_SUBSCRIBE, WS_TOPIC_UNSUBSCRIBE, ) +from ..models.config_data import ConfigData _LOGGER = logging.getLogger(__name__) -class IntegrationWS(BaseAPI): - _config_data: ConfigData | None +class WebSockets: + _hass: HomeAssistant | None + _session: ClientSession | None + _triggered_sensors: dict _api_data: dict - _can_log_messages: bool + _config_data: ConfigData + _entry_id: str | None + + _status: ConnectivityStatus | None + _on_status_changed: Callable[[ConnectivityStatus], Awaitable[None]] _previous_message: dict | None - _ws_handlers: dict def __init__( - self, - hass: HomeAssistant | None, - async_on_data_changed: Callable[[], Awaitable[None]] | None = None, - async_on_status_changed: Callable[[ConnectivityStatus], Awaitable[None]] - | None = None, + self, hass: HomeAssistant, config_data: ConfigData, entry_id: str | None = None ): - super().__init__(hass, async_on_data_changed, async_on_status_changed) + try: + self._hass = hass + self._config_data = config_data + self._entry_id = entry_id - self._config_data = None - self._ws = None - self._api_data = {} - self._remove_async_track_time = None - self._ws_handlers = self._get_ws_handlers() - self._can_log_messages: bool = False - self._previous_message = None + self._status = None + self._session = None + + self._base_url = None + self._pending_payloads = [] + self._ws = None + self._api_data = {} + self._data = { + WS_EXPORT_KEY: {}, + WS_INTERFACES_KEY: {}, + } + self._triggered_sensors = {} + self._remove_async_track_time = None + + self._local_async_dispatcher_send = None + + self._messages_handler: dict = self._get_ws_handlers() + + self._can_log_messages: bool = False + self._previous_message = None + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.error( + f"Failed to load MyDolphin Plus WS, error: {ex}, line: {line_number}" + ) + + @property + def data(self) -> dict: + return self._data + + @property + def status(self) -> str | None: + status = self._status + + return status + + @property + def _is_home_assistant(self): + return self._hass is not None + + @property + def _has_running_loop(self): + return self._hass.loop is not None and not self._hass.loop.is_closed() @property def _api_session_id(self): @@ -91,76 +138,46 @@ def _api_cookies(self): return api_cookies - @property - def _ws_url(self): - url = urlparse(self._config_data.url) - - ws_url = WEBSOCKET_URL_TEMPLATE.format(url.netloc) - - return ws_url - - async def update_api_data(self, api_data: dict, can_log_messages: bool): + def update_api_data(self, api_data: dict, can_log_messages: bool): self._api_data = api_data self._can_log_messages = can_log_messages - async def initialize(self, config_data: ConfigData | None = None): - if config_data is None: - _LOGGER.debug("Reinitializing WebSocket connection") - - else: - self._config_data = config_data - - _LOGGER.debug("Initializing WebSocket connection") - + async def initialize(self): try: - self.data = { - WS_EXPORT_KEY: {}, - WS_INTERFACES_KEY: {}, - } + _LOGGER.debug("Initializing") - await self.initialize_session(cookies=self._api_cookies) + await self._initialize_session() - async with self.session.ws_connect( - self._ws_url, + async with self._session.ws_connect( + self._config_data.ws_url, ssl=False, autoclose=True, max_msg_size=WS_MAX_MSG_SIZE, timeout=WS_TIMEOUT, compress=WS_COMPRESSION_DEFLATE, ) as ws: - await self.set_status(ConnectivityStatus.Connected) - self._ws = ws - await self._listen() - await self.set_status(ConnectivityStatus.NotConnected) - except Exception as ex: - if self.session is not None and self.session.closed: - _LOGGER.info("WS Session closed") - - await self.terminate() + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno - else: - exc_type, exc_obj, tb = sys.exc_info() - line_number = tb.tb_lineno + if self.status == ConnectivityStatus.Connected: + _LOGGER.info( + f"WS got disconnected will try to recover, Error: {ex}, Line: {line_number}" + ) - if self.status == ConnectivityStatus.Connected: - _LOGGER.info( - f"WS got disconnected will try to recover, Error: {ex}, Line: {line_number}" - ) + self._set_status(ConnectivityStatus.NotConnected) - else: - _LOGGER.warning( - f"Failed to connect WS, Error: {ex}, Line: {line_number}" - ) + else: + _LOGGER.warning( + f"Failed to connect WS, Error: {ex}, Line: {line_number}" + ) - await self.set_status(ConnectivityStatus.Failed) + self._set_status(ConnectivityStatus.Failed) async def terminate(self): - await super().terminate() - if self._remove_async_track_time is not None: self._remove_async_track_time() self._remove_async_track_time = None @@ -168,14 +185,32 @@ async def terminate(self): if self._ws is not None: await self._ws.close() + await asyncio.sleep(DISCONNECT_INTERVAL) + + self._set_status(ConnectivityStatus.Disconnected) self._ws = None - async def async_send_heartbeat(self): - _LOGGER.debug("Keep alive message sent") - if self.session is None or self.session.closed: - await self.set_status(ConnectivityStatus.NotConnected) + async def _initialize_session(self): + try: + if self._is_home_assistant: + self._session = async_create_clientsession(hass=self._hass) + + else: + self._session = ClientSession() + + except Exception as ex: + exc_type, exc_obj, tb = sys.exc_info() + line_number = tb.tb_lineno + + _LOGGER.warning( + f"Failed to initialize session, Error: {str(ex)}, Line: {line_number}" + ) - return + self._set_status(ConnectivityStatus.Failed) + + async def send_heartbeat(self): + if self._session is None or self._session.closed: + self._set_status(ConnectivityStatus.NotConnected) if self.status == ConnectivityStatus.Connected: content = {"CLIENT_PING": "", "SESSION_ID": self._api_session_id} @@ -193,7 +228,7 @@ async def async_send_heartbeat(self): _LOGGER.debug( f"Gracefully failed to send heartbeat - Restarting connection, Error: {crex}" ) - await self.set_status(ConnectivityStatus.NotConnected) + self._set_status(ConnectivityStatus.NotConnected) except Exception as ex: _LOGGER.error(f"Failed to send heartbeat, Error: {ex}") @@ -204,39 +239,50 @@ async def _listen(self): subscription_data = self._get_subscription_data() await self._ws.send_str(subscription_data) - _LOGGER.info("Subscribed to WS payloads") + self._set_status(ConnectivityStatus.Connected) async for msg in self._ws: + is_ha_running = self._hass.is_running is_connected = self.status == ConnectivityStatus.Connected is_closing_type = msg.type in WS_CLOSING_MESSAGE is_error = msg.type == aiohttp.WSMsgType.ERROR + can_try_parse_message = msg.type == aiohttp.WSMsgType.TEXT is_closing_data = ( False if is_closing_type or is_error else msg.data == "close" ) - session_is_closed = self.session is None or self.session.closed - - if ( - is_closing_type - or is_error - or is_closing_data - or session_is_closed - or not is_connected - ): + session_is_closed = self._session is None or self._session.closed + + not_connected = True in [ + is_closing_type, + is_error, + is_closing_data, + session_is_closed, + not is_connected, + ] + + if not is_ha_running: + self._set_status(ConnectivityStatus.Disconnected) + return + + if not_connected: _LOGGER.warning( f"WS stopped listening, " f"Message: {str(msg)}, " f"Exception: {self._ws.exception()}" ) - break + self._set_status(ConnectivityStatus.NotConnected) + return - else: + elif can_try_parse_message: if self._can_log_messages: _LOGGER.debug(f"Message received: {str(msg)}") - await self.parse_message(msg.data) + self.data[API_DATA_LAST_UPDATE] = datetime.now().isoformat() + + await self._parse_message(msg.data) - async def parse_message(self, message): + async def _parse_message(self, message: str): try: self._increase_counter(WS_RECEIVED_MESSAGES) @@ -250,6 +296,8 @@ async def parse_message(self, message): payload_json = json.loads(message_json) await self._message_handler(payload_json) + + self._async_dispatcher_send(SIGNAL_DATA_CHANGED) else: self._increase_counter(WS_IGNORED_MESSAGES) @@ -311,7 +359,7 @@ def _get_corrected_message(self, message): return message def _get_subscription_data(self): - topics = self._ws_handlers.keys() + topics = self._messages_handler.keys() topics_to_subscribe = [{WS_TOPIC_NAME: topic} for topic in topics] topics_to_unsubscribe = [] @@ -331,6 +379,42 @@ def _get_subscription_data(self): return data + def _set_status(self, status: ConnectivityStatus): + if status != self._status: + log_level = ConnectivityStatus.get_log_level(status) + + _LOGGER.log( + log_level, + f"Status changed from '{self._status}' to '{status}'", + ) + + self._status = status + + self._async_dispatcher_send( + SIGNAL_WS_STATUS, + status, + ) + + def set_local_async_dispatcher_send(self, dispatcher_send): + self._local_async_dispatcher_send = dispatcher_send + + def _async_dispatcher_send(self, signal: str, *args: Any) -> None: + if self._hass is None: + self._local_async_dispatcher_send(signal, self._entry_id, *args) + + else: + async_dispatcher_send(self._hass, signal, self._entry_id, *args) + + def _increase_counter(self, key): + counter = self.data.get(key, 0) + + self.data[key] = counter + 1 + + def _decrease_counter(self, key): + counter = self.data.get(key, 0) + + self.data[key] = counter - 1 + def _get_ws_handlers(self) -> dict: ws_handlers = { WS_EXPORT_KEY: self._handle_export, @@ -346,14 +430,13 @@ async def _message_handler(self, payload=None): if payload is not None: for key in payload: data = payload.get(key) - handler = self._ws_handlers.get(key) + handler = self._messages_handler.get(key) if handler is None: _LOGGER.error(f"Handler not found for {key}") else: handler(data) - await self.fire_data_changed_event() except Exception as ex: exc_type, exc_obj, tb = sys.exc_info() line_number = tb.tb_lineno @@ -501,13 +584,3 @@ def _handle_discover(self, data): _LOGGER.error( f"Failed to load {WS_DISCOVER_KEY}, Original Message: {data}, Error: {ex}, Line: {line_number}" ) - - def _increase_counter(self, key): - counter = self.data.get(key, 0) - - self.data[key] = counter + 1 - - def _decrease_counter(self, key): - counter = self.data.get(key, 0) - - self.data[key] = counter - 1 diff --git a/custom_components/edgeos/manifest.json b/custom_components/edgeos/manifest.json index 8581079..b6fa548 100644 --- a/custom_components/edgeos/manifest.json +++ b/custom_components/edgeos/manifest.json @@ -8,5 +8,5 @@ "iot_class": "local_polling", "issue_tracker": "https://github.com/elad-bar/ha-edgeos/issues", "requirements": ["aiohttp"], - "version": "2.0.32" + "version": "2.1.0" } diff --git a/custom_components/edgeos/models/config_data.py b/custom_components/edgeos/models/config_data.py new file mode 100644 index 0000000..e50f515 --- /dev/null +++ b/custom_components/edgeos/models/config_data.py @@ -0,0 +1,101 @@ +from urllib.parse import urlparse + +import voluptuous as vol +from voluptuous import Schema + +from homeassistant.const import ( + CONF_HOST, + CONF_PASSWORD, + CONF_PATH, + CONF_PORT, + CONF_SSL, + CONF_USERNAME, +) + +from ..common.consts import ( + API_URL_TEMPLATE, + CONF_TITLE, + DEFAULT_NAME, + WEBSOCKET_URL_TEMPLATE, +) + +DATA_KEYS = [CONF_HOST, CONF_PORT, CONF_SSL, CONF_PATH, CONF_USERNAME, CONF_PASSWORD] + + +class ConfigData: + _hostname: str | None + _username: str | None + _password: str | None + + def __init__(self): + self._hostname = None + self._username = None + self._password = None + + @property + def hostname(self) -> str: + hostname = self._hostname + + return hostname + + @property + def username(self) -> str: + username = self._username + + return username + + @property + def password(self) -> str: + password = self._password + + return password + + @property + def api_url(self): + url = API_URL_TEMPLATE.format(self._hostname) + + return url + + @property + def ws_url(self): + url = urlparse(self.api_url) + + ws_url = WEBSOCKET_URL_TEMPLATE.format(url.netloc) + + return ws_url + + def update(self, data: dict): + self._password = data.get(CONF_PASSWORD) + self._username = data.get(CONF_USERNAME) + self._hostname = data.get(CONF_HOST) + + def to_dict(self): + obj = { + CONF_USERNAME: self.username, + CONF_HOST: self.hostname, + } + + return obj + + def __repr__(self): + to_string = f"{self.to_dict()}" + + return to_string + + @staticmethod + def default_schema(user_input: dict | None) -> Schema: + if user_input is None: + user_input = {} + + new_user_input = { + vol.Required( + CONF_TITLE, default=user_input.get(CONF_TITLE, DEFAULT_NAME) + ): str, + vol.Required(CONF_HOST, default=user_input.get(CONF_HOST)): str, + vol.Required(CONF_USERNAME, default=user_input.get(CONF_USERNAME)): str, + vol.Required(CONF_PASSWORD, default=user_input.get(CONF_PASSWORD)): str, + } + + schema = vol.Schema(new_user_input) + + return schema diff --git a/custom_components/edgeos/component/models/edge_os_device_data.py b/custom_components/edgeos/models/edge_os_device_data.py similarity index 87% rename from custom_components/edgeos/component/models/edge_os_device_data.py rename to custom_components/edgeos/models/edge_os_device_data.py index 7150b41..f2d406b 100644 --- a/custom_components/edgeos/component/models/edge_os_device_data.py +++ b/custom_components/edgeos/models/edge_os_device_data.py @@ -2,8 +2,7 @@ from datetime import datetime, timedelta -from ...core.helpers.const import ENTITY_UNIQUE_ID -from ..helpers.const import ( +from ..common.consts import ( DEVICE_DATA_DOMAIN, DEVICE_DATA_IP, DEVICE_DATA_MAC, @@ -82,12 +81,22 @@ def to_dict(self): DEVICE_DATA_DOMAIN: self.domain, DEVICE_DATA_RECEIVED: self.received.to_dict(), DEVICE_DATA_SENT: self.sent.to_dict(), - ENTITY_UNIQUE_ID: self.unique_id, DHCP_SERVER_LEASED: self.is_leased, } return obj + def get_attributes(self): + device_attributes = self.to_dict() + + attributes = { + attribute: device_attributes[attribute] + for attribute in device_attributes + if attribute not in [DEVICE_DATA_RECEIVED, DEVICE_DATA_SENT] + } + + return attributes + def __repr__(self): to_string = f"{self.to_dict()}" diff --git a/custom_components/edgeos/component/models/edge_os_interface_data.py b/custom_components/edgeos/models/edge_os_interface_data.py similarity index 74% rename from custom_components/edgeos/component/models/edge_os_interface_data.py rename to custom_components/edgeos/models/edge_os_interface_data.py index 4116c18..fe565ba 100644 --- a/custom_components/edgeos/component/models/edge_os_interface_data.py +++ b/custom_components/edgeos/models/edge_os_interface_data.py @@ -1,16 +1,14 @@ from __future__ import annotations -from ...core.helpers.const import ENTITY_UNIQUE_ID -from ..helpers.const import ( - IGNORED_INTERFACES, +from ..common.consts import ( INTERFACE_DATA_ADDRESS, INTERFACE_DATA_AGING, INTERFACE_DATA_BRIDGE_GROUP, INTERFACE_DATA_BRIDGED_CONNTRACK, INTERFACE_DATA_DESCRIPTION, INTERFACE_DATA_DUPLEX, - INTERFACE_DATA_HANDLER, INTERFACE_DATA_HELLO_TIME, + INTERFACE_DATA_IS_SUPPORTED, INTERFACE_DATA_MAX_AGE, INTERFACE_DATA_MULTICAST, INTERFACE_DATA_NAME, @@ -21,6 +19,7 @@ INTERFACE_DATA_SPEED, INTERFACE_DATA_STP, INTERFACE_DATA_TYPE, + INTERFACE_DYNAMIC_SUPPORTED, RECEIVED_DROPPED_PREFIX, RECEIVED_ERRORS_PREFIX, RECEIVED_PACKETS_PREFIX, @@ -31,17 +30,16 @@ SENT_PACKETS_PREFIX, SENT_RATE_PREFIX, SENT_TRAFFIC_PREFIX, - SPECIAL_INTERFACES, TRAFFIC_DATA_DIRECTION_RECEIVED, TRAFFIC_DATA_DIRECTION_SENT, ) -from ..helpers.enums import InterfaceHandlers +from ..common.enums import DynamicInterfaceTypes, InterfaceTypes from .edge_os_traffic_data import EdgeOSTrafficData class EdgeOSInterfaceData: name: str - interface_type: str | None + interface_type: InterfaceTypes | None duplex: str | None speed: str | None description: str | None @@ -60,11 +58,11 @@ class EdgeOSInterfaceData: up: bool | None l1up: bool | None mac: str | None - handler: InterfaceHandlers + is_supported: bool - def __init__(self, name: str): + def __init__(self, name: str, interface_type: InterfaceTypes): self.name = name - self.interface_type = None + self.interface_type = interface_type self.description = None self.duplex = None self.speed = None @@ -84,7 +82,7 @@ def __init__(self, name: str): self.up = None self.l1up = None self.mac = None - self.handler = InterfaceHandlers.IGNORED + self.is_supported = self._get_is_supported() @property def unique_id(self) -> str: @@ -95,7 +93,7 @@ def to_dict(self): INTERFACE_DATA_NAME: self.name, INTERFACE_DATA_DESCRIPTION: self.description, INTERFACE_DATA_TYPE: self.interface_type, - INTERFACE_DATA_HANDLER: self.handler.name, + INTERFACE_DATA_IS_SUPPORTED: self.is_supported, INTERFACE_DATA_DUPLEX: self.duplex, INTERFACE_DATA_SPEED: self.speed, INTERFACE_DATA_BRIDGE_GROUP: self.bridge_group, @@ -110,28 +108,25 @@ def to_dict(self): INTERFACE_DATA_MULTICAST: self.multicast, INTERFACE_DATA_RECEIVED: self.received.to_dict(), INTERFACE_DATA_SENT: self.sent.to_dict(), - ENTITY_UNIQUE_ID: self.unique_id, } return obj - def set_type(self, interface_type: str | None): - handler = InterfaceHandlers.IGNORED + def _get_is_supported(self): + is_supported = self.interface_type in [ + InterfaceTypes.ETHERNET, + InterfaceTypes.BRIDGE, + ] - if interface_type is None: - for special_interface in SPECIAL_INTERFACES: - if self.name.startswith(special_interface): - handler = InterfaceHandlers.SPECIAL - interface_type = SPECIAL_INTERFACES.get(special_interface) + if self.interface_type == InterfaceTypes.DYNAMIC: + prefixes = list(DynamicInterfaceTypes) + for prefix in prefixes: + if self.name.startswith(str(prefix)): + is_supported = prefix in INTERFACE_DYNAMIC_SUPPORTED break - else: - if interface_type not in IGNORED_INTERFACES: - handler = InterfaceHandlers.REGULAR - - self.handler = handler - self.interface_type = interface_type + return is_supported def get_stats(self): data = { @@ -149,6 +144,23 @@ def get_stats(self): return data + def get_attributes(self): + interface_attributes = self.to_dict() + + attributes = { + attribute: interface_attributes[attribute] + for attribute in interface_attributes + if attribute + not in [ + INTERFACE_DATA_RECEIVED, + INTERFACE_DATA_SENT, + INTERFACE_DATA_IS_SUPPORTED, + ] + and interface_attributes[attribute] is not None + } + + return attributes + def __repr__(self): to_string = f"{self.to_dict()}" diff --git a/custom_components/edgeos/component/models/edge_os_system_data.py b/custom_components/edgeos/models/edge_os_system_data.py similarity index 98% rename from custom_components/edgeos/component/models/edge_os_system_data.py rename to custom_components/edgeos/models/edge_os_system_data.py index 5fe3453..c9db133 100644 --- a/custom_components/edgeos/component/models/edge_os_system_data.py +++ b/custom_components/edgeos/models/edge_os_system_data.py @@ -2,7 +2,7 @@ from datetime import datetime -from ..helpers.const import ( +from ..common.consts import ( DHCP_SERVER_LEASES, SYSTEM_DATA_DISABLE, SYSTEM_DATA_ENABLE, diff --git a/custom_components/edgeos/component/models/edge_os_traffic_data.py b/custom_components/edgeos/models/edge_os_traffic_data.py similarity index 96% rename from custom_components/edgeos/component/models/edge_os_traffic_data.py rename to custom_components/edgeos/models/edge_os_traffic_data.py index d3a5770..fdb30c9 100644 --- a/custom_components/edgeos/component/models/edge_os_traffic_data.py +++ b/custom_components/edgeos/models/edge_os_traffic_data.py @@ -2,7 +2,7 @@ from datetime import datetime, timedelta -from ..helpers.const import ( +from ..common.consts import ( TRAFFIC_DATA_DIRECTION, TRAFFIC_DATA_DROPPED, TRAFFIC_DATA_ERRORS, @@ -41,7 +41,7 @@ def update(self, data: dict): if self.rate > 0: now = datetime.now().timestamp() - self.last_activity = int(now) + self.last_activity = float(now) def to_dict(self): now = datetime.now().timestamp() diff --git a/custom_components/edgeos/models/exceptions.py b/custom_components/edgeos/models/exceptions.py new file mode 100644 index 0000000..05e0b03 --- /dev/null +++ b/custom_components/edgeos/models/exceptions.py @@ -0,0 +1,49 @@ +from homeassistant.exceptions import HomeAssistantError + +from ..common.connectivity_status import ConnectivityStatus + + +class LoginError(Exception): + def __init__(self): + self.error = "Failed to login" + + +class AlreadyExistsError(HomeAssistantError): + title: str + + def __init__(self, title: str): + self.title = title + + +class APIValidationException(HomeAssistantError): + endpoint: str + status: ConnectivityStatus + + def __init__(self, endpoint: str, status: ConnectivityStatus): + super().__init__( + f"API cannot process request to '{endpoint}', Status: {status}" + ) + + self.endpoint = endpoint + self.status = status + + +class IncompatibleVersion(HomeAssistantError): + def __init__(self, version): + self._version = version + + def __repr__(self): + return f"Unsupported EdgeOS version ({self._version})" + + +class SessionTerminatedException(HomeAssistantError): + Terminated = True + + +class LoginException(HomeAssistantError): + def __init__(self, status_code): + self._status_code = status_code + + @property + def status_code(self): + return self._status_code diff --git a/custom_components/edgeos/number.py b/custom_components/edgeos/number.py new file mode 100644 index 0000000..ca60995 --- /dev/null +++ b/custom_components/edgeos/number.py @@ -0,0 +1,65 @@ +from abc import ABC +import logging + +from homeassistant.components.number import NumberEntity +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ATTR_STATE, Platform +from homeassistant.core import HomeAssistant + +from .common.base_entity import IntegrationBaseEntity, async_setup_base_entry +from .common.consts import ACTION_ENTITY_SET_NATIVE_VALUE, ATTR_ATTRIBUTES +from .common.entity_descriptions import IntegrationNumberEntityDescription +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator + +_LOGGER = logging.getLogger(__name__) + + +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities +): + await async_setup_base_entry( + hass, + entry, + Platform.NUMBER, + IntegrationNumberEntity, + async_add_entities, + ) + + +class IntegrationNumberEntity(IntegrationBaseEntity, NumberEntity, ABC): + """Representation of a sensor.""" + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationNumberEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(hass, entity_description, coordinator, device_type, item_id) + + self.entity_description = entity_description + + self._attr_native_min_value = entity_description.native_min_value + self._attr_native_max_value = entity_description.native_max_value + self._attr_native_step = 1 + + async def async_set_native_value(self, value: float) -> None: + """Change the selected option.""" + await self.async_execute_device_action( + ACTION_ENTITY_SET_NATIVE_VALUE, int(value) + ) + + def update_component(self, data): + """Fetch new state parameters for the sensor.""" + if data is not None: + state = data.get(ATTR_STATE) + attributes = data.get(ATTR_ATTRIBUTES) + + self._attr_native_value = int(state) + self._attr_extra_state_attributes = attributes + + else: + self._attr_native_value = None diff --git a/custom_components/edgeos/select.py b/custom_components/edgeos/select.py index facfc33..84a7cb1 100644 --- a/custom_components/edgeos/select.py +++ b/custom_components/edgeos/select.py @@ -1,32 +1,62 @@ -""" -Support for select. -""" -from __future__ import annotations - +from abc import ABC import logging -from .core.components.select import CoreSelect -from .core.helpers.setup_base_entry import async_setup_base_entry +from homeassistant.components.select import SelectEntity +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ATTR_STATE, Platform +from homeassistant.core import HomeAssistant + +from .common.base_entity import IntegrationBaseEntity, async_setup_base_entry +from .common.consts import ACTION_ENTITY_SELECT_OPTION, ATTR_ATTRIBUTES +from .common.entity_descriptions import IntegrationSelectEntityDescription +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator _LOGGER = logging.getLogger(__name__) -async def async_setup_entry(hass, config_entry, async_add_devices): - """Set up the component.""" +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities +): await async_setup_base_entry( hass, - config_entry, - async_add_devices, - CoreSelect.get_domain(), - CoreSelect.get_component, + entry, + Platform.SELECT, + IntegrationSelectEntity, + async_add_entities, ) -async def async_unload_entry(hass, config_entry): - _LOGGER.info(f"Unload entry for {CoreSelect.get_domain()} domain: {config_entry}") +class IntegrationSelectEntity(IntegrationBaseEntity, SelectEntity, ABC): + """Representation of a sensor.""" + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationSelectEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(hass, entity_description, coordinator, device_type, item_id) + + self.entity_description = entity_description + + self._attr_options = entity_description.options + self._attr_current_option = entity_description.options[0] + + async def async_select_option(self, option: str) -> None: + """Change the selected option.""" + await self.async_execute_device_action(ACTION_ENTITY_SELECT_OPTION, option) - return True + def update_component(self, data): + """Fetch new state parameters for the sensor.""" + if data is not None: + state = data.get(ATTR_STATE) + attributes = data.get(ATTR_ATTRIBUTES) + self._attr_current_option = state + self._attr_extra_state_attributes = attributes -async def async_remove_entry(hass, entry) -> None: - _LOGGER.info(f"Remove entry for {CoreSelect.get_domain()} entry: {entry}") + else: + self._attr_current_option = self.entity_description.options[0] diff --git a/custom_components/edgeos/sensor.py b/custom_components/edgeos/sensor.py index f57921a..ca3b24a 100644 --- a/custom_components/edgeos/sensor.py +++ b/custom_components/edgeos/sensor.py @@ -1,32 +1,58 @@ -""" -Support for sensor -""" -from __future__ import annotations - import logging -from .core.components.sensor import CoreSensor -from .core.helpers.setup_base_entry import async_setup_base_entry +from homeassistant.components.sensor import SensorEntity +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ATTR_ICON, ATTR_STATE, Platform +from homeassistant.core import HomeAssistant + +from .common.base_entity import IntegrationBaseEntity, async_setup_base_entry +from .common.consts import ATTR_ATTRIBUTES +from .common.entity_descriptions import IntegrationSensorEntityDescription +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator _LOGGER = logging.getLogger(__name__) -async def async_setup_entry(hass, config_entry, async_add_devices): - """Set up the Sensor.""" +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities +): await async_setup_base_entry( hass, - config_entry, - async_add_devices, - CoreSensor.get_domain(), - CoreSensor.get_component, + entry, + Platform.SENSOR, + IntegrationSensorEntity, + async_add_entities, ) -async def async_unload_entry(hass, config_entry): - _LOGGER.info(f"Unload entry for {CoreSensor.get_domain()} domain: {config_entry}") +class IntegrationSensorEntity(IntegrationBaseEntity, SensorEntity): + """Representation of a sensor.""" + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationSensorEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(hass, entity_description, coordinator, device_type, item_id) + + self._attr_device_class = entity_description.device_class + + def update_component(self, data): + """Fetch new state parameters for the sensor.""" + if data is not None: + state = data.get(ATTR_STATE) + attributes = data.get(ATTR_ATTRIBUTES) + icon = data.get(ATTR_ICON) - return True + self._attr_native_value = state + self._attr_extra_state_attributes = attributes + if icon is not None: + self._attr_icon = icon -async def async_remove_entry(hass, entry) -> None: - _LOGGER.info(f"Remove entry for {CoreSensor.get_domain()} entry: {entry}") + else: + self._attr_native_value = None diff --git a/custom_components/edgeos/strings.json b/custom_components/edgeos/strings.json index c4d8dff..a531de2 100644 --- a/custom_components/edgeos/strings.json +++ b/custom_components/edgeos/strings.json @@ -48,5 +48,101 @@ "already_configured": "EdgeOS integration already configured with the name", "incompatible_version": "Unsupported firmware version" } + }, + "entity": { + "binary_sensor": { + "firmware": { + "name": "Firmware Upgrade" + }, + "interface_connected": { + "name": "Connected" + } + }, + "device_tracker": { + "device_tracker": { + "name": "Status" + } + }, + "number": { + "consider_away_interval": { + "name": "Consider Away Interval" + }, + "update_entities_interval": { + "name": "Update Entities Interval" + }, + "update_api_interval": { + "name": "Update API Interval" + } + }, + "sensor": { + "last_restart": { + "name": "Last Restart" + }, + "cpu_usage": { + "name": "CPU Usage" + }, + "ram_usage": { + "name": "RAM Usage" + }, + "unknown_devices": { + "name": "Unknown Devices" + }, + "interface_received_dropped": { + "name": "Received Dropped" + }, + "interface_sent_dropped": { + "name": "Sent Dropped" + }, + "interface_received_errors": { + "name": "Received Errors" + }, + "interface_sent_errors": { + "name": "Sent Errors" + }, + "interface_received_packets": { + "name": "Received Packets" + }, + "interface_sent_packets": { + "name": "Sent Packets" + }, + "interface_received_rate": { + "name": "Received Rate" + }, + "interface_sent_rate": { + "name": "Sent Rate" + }, + "interface_received_traffic": { + "name": "Received Traffic" + }, + "interface_sent_traffic": { + "name": "Sent Traffic" + }, + "device_received_rate": { + "name": "Received Rate" + }, + "device_sent_rate": { + "name": "Sent Rate" + }, + "device_received_traffic": { + "name": "Received Traffic" + }, + "device_sent_traffic": { + "name": "Sent Traffic" + } + }, + "switch": { + "log_incoming_messages": { + "name": "Log Incoming Messages" + }, + "interface_monitored": { + "name": "Monitored" + }, + "interface_status": { + "name": "Status" + }, + "device_monitored": { + "name": "Monitored" + } + } } } diff --git a/custom_components/edgeos/switch.py b/custom_components/edgeos/switch.py index 5ef5c99..fb788c8 100644 --- a/custom_components/edgeos/switch.py +++ b/custom_components/edgeos/switch.py @@ -1,32 +1,73 @@ -""" -Support for switch. -""" -from __future__ import annotations - +from abc import ABC import logging +from typing import Any + +from homeassistant.components.switch import SwitchEntity +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ATTR_ICON, Platform +from homeassistant.core import HomeAssistant -from .core.components.switch import CoreSwitch -from .core.helpers.setup_base_entry import async_setup_base_entry +from .common.base_entity import IntegrationBaseEntity, async_setup_base_entry +from .common.consts import ( + ACTION_ENTITY_TURN_OFF, + ACTION_ENTITY_TURN_ON, + ATTR_ATTRIBUTES, + ATTR_IS_ON, +) +from .common.entity_descriptions import IntegrationSwitchEntityDescription +from .common.enums import DeviceTypes +from .managers.coordinator import Coordinator _LOGGER = logging.getLogger(__name__) -async def async_setup_entry(hass, config_entry, async_add_devices): - """Set up the Switch.""" +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities +): await async_setup_base_entry( hass, - config_entry, - async_add_devices, - CoreSwitch.get_domain(), - CoreSwitch.get_component, + entry, + Platform.SWITCH, + IntegrationSwitchEntity, + async_add_entities, ) -async def async_unload_entry(hass, config_entry): - _LOGGER.info(f"Unload entry for {CoreSwitch.get_domain()} domain: {config_entry}") +class IntegrationSwitchEntity(IntegrationBaseEntity, SwitchEntity, ABC): + """Representation of a sensor.""" + + def __init__( + self, + hass: HomeAssistant, + entity_description: IntegrationSwitchEntityDescription, + coordinator: Coordinator, + device_type: DeviceTypes, + item_id: str | None, + ): + super().__init__(hass, entity_description, coordinator, device_type, item_id) + + self._attr_device_class = entity_description.device_class + + async def async_turn_on(self, **kwargs: Any) -> None: + """Turn the entity on.""" + await self.async_execute_device_action(ACTION_ENTITY_TURN_ON) + + async def async_turn_off(self, **kwargs: Any) -> None: + """Turn the entity off.""" + await self.async_execute_device_action(ACTION_ENTITY_TURN_OFF) + + def update_component(self, data): + """Fetch new state parameters for the sensor.""" + if data is not None: + is_on = data.get(ATTR_IS_ON) + attributes = data.get(ATTR_ATTRIBUTES) + icon = data.get(ATTR_ICON) - return True + self._attr_is_on = is_on + self._attr_extra_state_attributes = attributes + if icon is not None: + self._attr_icon = icon -async def async_remove_entry(hass, entry) -> None: - _LOGGER.info(f"Remove entry for {CoreSwitch.get_domain()} entry: {entry}") + else: + self._attr_is_on = None diff --git a/custom_components/edgeos/translations/en.json b/custom_components/edgeos/translations/en.json index 3da7b8d..04a182f 100644 --- a/custom_components/edgeos/translations/en.json +++ b/custom_components/edgeos/translations/en.json @@ -1,51 +1,148 @@ { "config": { + "abort": { + "already_configured": "EdgeOS integration ({name}) already configured" + }, + "error": { + "already_configured": "EdgeOS integration already configured with the name", + "auth_general_error": "General failure, please try again", + "empty_device_data": "Could not retrieve device data from EdgeOS Router", + "incompatible_version": "Unsupported firmware version", + "invalid_credentials": "Invalid credentials to EdgeOS Router", + "invalid_dpi_configuration": "Deep Packet Inspection (traffic-analysis) configuration is disabled, please enable", + "invalid_export_configuration": "Export (traffic-analysis) configuration is disabled, please enable", + "not_found": "EdgeOS URL not found" + }, "step": { "user": { - "title": "Set up EdgeOS Router", - "description": "Set up your EdgeOS Router, note that integration will communicate using HTTPS connection", "data": { "host": "Host", - "username": "Username", - "password": "Password" - } + "password": "Password", + "username": "Username" + }, + "description": "Set up your EdgeOS Router, note that integration will communicate using HTTPS connection", + "title": "Set up EdgeOS Router" + } + } + }, + "entity": { + "binary_sensor": { + "firmware": { + "name": "Firmware Upgrade" + }, + "interface_connected": { + "name": "Connected" } }, - "abort": { - "already_configured": "EdgeOS integration ({name}) already configured" + "device_tracker": { + "device_tracker": { + "name": "Status" + } }, - "error": { - "invalid_credentials": "Invalid credentials to EdgeOS Router", - "auth_general_error": "General failure, please try again", - "not_found": "EdgeOS URL not found", - "invalid_export_configuration": "Export (traffic-analysis) configuration is disabled, please enable", - "invalid_dpi_configuration": "Deep Packet Inspection (traffic-analysis) configuration is disabled, please enable", - "empty_device_data": "Could not retrieve device data from EdgeOS Router", - "already_configured": "EdgeOS integration already configured with the name", - "incompatible_version": "Unsupported firmware version" + "number": { + "consider_away_interval": { + "name": "Consider Away Interval" + }, + "update_api_interval": { + "name": "Update API Interval" + }, + "update_entities_interval": { + "name": "Update Entities Interval" + } + }, + "sensor": { + "cpu_usage": { + "name": "CPU Usage" + }, + "device_received_rate": { + "name": "Received Rate" + }, + "device_received_traffic": { + "name": "Received Traffic" + }, + "device_sent_rate": { + "name": "Sent Rate" + }, + "device_sent_traffic": { + "name": "Sent Traffic" + }, + "interface_received_dropped": { + "name": "Received Dropped" + }, + "interface_received_errors": { + "name": "Received Errors" + }, + "interface_received_packets": { + "name": "Received Packets" + }, + "interface_received_rate": { + "name": "Received Rate" + }, + "interface_received_traffic": { + "name": "Received Traffic" + }, + "interface_sent_dropped": { + "name": "Sent Dropped" + }, + "interface_sent_errors": { + "name": "Sent Errors" + }, + "interface_sent_packets": { + "name": "Sent Packets" + }, + "interface_sent_rate": { + "name": "Sent Rate" + }, + "interface_sent_traffic": { + "name": "Sent Traffic" + }, + "last_restart": { + "name": "Last Restart" + }, + "ram_usage": { + "name": "RAM Usage" + }, + "unknown_devices": { + "name": "Unknown Devices" + } + }, + "switch": { + "device_monitored": { + "name": "Monitored" + }, + "interface_monitored": { + "name": "Monitored" + }, + "interface_status": { + "name": "Status" + }, + "log_incoming_messages": { + "name": "Log Incoming Messages" + } } }, "options": { + "error": { + "already_configured": "EdgeOS integration already configured with the name", + "auth_general_error": "General failure, please try again", + "empty_device_data": "Could not retrieve device data from EdgeOS Router", + "incompatible_version": "Unsupported firmware version", + "invalid_credentials": "Invalid credentials to EdgeOS Router", + "invalid_dpi_configuration": "Deep Packet Inspection (traffic-analysis) configuration is disabled, please enable", + "invalid_export_configuration": "Export (traffic-analysis) configuration is disabled, please enable", + "not_found": "EdgeOS URL not found" + }, "step": { "edge_os_additional_settings": { - "title": "Options for EdgeOS.", - "description": "Set up your monitored devices / interfaces and tracked devices, use comma separated values, to clear existing value, use the checkbox.", "data": { + "clear-credentials": "Clear credentials", "host": "Host", - "username": "Username", - "password": "Password" - } + "password": "Password", + "username": "Username" + }, + "description": "Set up your monitored devices / interfaces and tracked devices, use comma separated values, to clear existing value, use the checkbox.", + "title": "Options for EdgeOS." } - }, - "error": { - "invalid_credentials": "Invalid credentials to EdgeOS Router", - "auth_general_error": "General failure, please try again", - "not_found": "EdgeOS URL not found", - "invalid_export_configuration": "Export (traffic-analysis) configuration is disabled, please enable", - "invalid_dpi_configuration": "Deep Packet Inspection (traffic-analysis) configuration is disabled, please enable", - "empty_device_data": "Could not retrieve device data from EdgeOS Router", - "already_configured": "EdgeOS integration already configured with the name", - "incompatible_version": "Unsupported firmware version" } } } diff --git a/custom_components/edgeos/translations/nb.json b/custom_components/edgeos/translations/nb.json index 7c5612f..86056f5 100644 --- a/custom_components/edgeos/translations/nb.json +++ b/custom_components/edgeos/translations/nb.json @@ -1,52 +1,148 @@ { "config": { + "abort": { + "already_configured": "Edgeos integrasjon ({name}) allerede konfigurert" + }, + "error": { + "already_configured": "Edgeos integrasjon allerede konfigurert med navnet", + "auth_general_error": "Generell fiasko, pr\u00f8v igjen", + "empty_device_data": "Kunne ikke hente enhetsdata fra Edgeos -ruteren", + "incompatible_version": "Ikke st\u00f8ttet firmwareversjon", + "invalid_credentials": "Ugyldig legitimasjon til Edgeos -ruteren", + "invalid_dpi_configuration": "Dyp pakkeinspeksjon (trafikkanalyse) Konfigurasjon er deaktivert, vennligst aktiver", + "invalid_export_configuration": "Eksport (trafikkanalyse) Konfigurasjon er deaktivert, vennligst aktiver", + "not_found": "Edgeos url ikke funnet" + }, "step": { "user": { - "title": "Sett opp EdgeOS Router", - "description": "Sett opp EdgeOS-ruteren din, vær oppmerksom på at integrasjon vil kommunisere ved hjelp av HTTPS-tilkobling", "data": { "host": "Vert", - "username": "Brukernavn", - "password": "Passord" - } + "password": "Passord", + "username": "Brukernavn" + }, + "description": "Sett opp Edgeos -ruteren, merk at integrasjon vil kommunisere ved hjelp av HTTPS -tilkobling", + "title": "Sett opp Edgeos -ruteren" + } + } + }, + "entity": { + "binary_sensor": { + "firmware": { + "name": "Firmwareoppgradering" + }, + "interface_connected": { + "name": "Tilkoblet" } }, - "abort": { - "already_configured": "EdgeOS-integrasjon ({name}) er allerede konfigurert" + "device_tracker": { + "device_tracker": { + "name": "Status" + } }, - "error": { - "invalid_credentials": "Ugyldig legitimasjon til EdgeOS Router", - "auth_general_error": "Generell feil, prøv igjen", - "not_found": "EdgeOS URL ikke funne", - "invalid_export_configuration": "Eksport (trafikkanalyse) konfigurasjon er deaktivert, vennligst aktiver", - "invalid_dpi_configuration": "Konfigurasjon av dyp pakkeinspeksjon (trafikkanalyse) er deaktivert, vennligst aktiver", - "empty_device_data": "Kunne ikke hente enhetsdata fra EdgeOS Router", - "already_configured": "EdgeOS-integrasjon allerede konfigurert med navnet", - "incompatible_version": "Inkompatibel versjon (påkrevd minst v2.0)" + "number": { + "consider_away_interval": { + "name": "Vurder bort intervall" + }, + "update_api_interval": { + "name": "Oppdater API -intervall" + }, + "update_entities_interval": { + "name": "Oppdater enheter intervall" + } + }, + "sensor": { + "cpu_usage": { + "name": "CPU bruk" + }, + "device_received_rate": { + "name": "Mottatt rente" + }, + "device_received_traffic": { + "name": "Mottatt trafikk" + }, + "device_sent_rate": { + "name": "Sendt rente" + }, + "device_sent_traffic": { + "name": "Sendte trafikk" + }, + "interface_received_dropped": { + "name": "Mottatt droppet" + }, + "interface_received_errors": { + "name": "Mottok feil" + }, + "interface_received_packets": { + "name": "Mottatte pakker" + }, + "interface_received_rate": { + "name": "Mottatt rente" + }, + "interface_received_traffic": { + "name": "Mottatt trafikk" + }, + "interface_sent_dropped": { + "name": "Sendt droppet" + }, + "interface_sent_errors": { + "name": "Sendte feil" + }, + "interface_sent_packets": { + "name": "Sendte pakker" + }, + "interface_sent_rate": { + "name": "Sendt rente" + }, + "interface_sent_traffic": { + "name": "Sendte trafikk" + }, + "last_restart": { + "name": "Siste omstart" + }, + "ram_usage": { + "name": "RAM -bruk" + }, + "unknown_devices": { + "name": "Ukjente enheter" + } + }, + "switch": { + "device_monitored": { + "name": "Overv\u00e5ket" + }, + "interface_monitored": { + "name": "Overv\u00e5ket" + }, + "interface_status": { + "name": "Status" + }, + "log_incoming_messages": { + "name": "Logg innkommende meldinger" + } } }, "options": { + "error": { + "already_configured": "Edgeos integrasjon allerede konfigurert med navnet", + "auth_general_error": "Generell fiasko, pr\u00f8v igjen", + "empty_device_data": "Kunne ikke hente enhetsdata fra Edgeos -ruteren", + "incompatible_version": "Ikke st\u00f8ttet firmwareversjon", + "invalid_credentials": "Ugyldig legitimasjon til Edgeos -ruteren", + "invalid_dpi_configuration": "Dyp pakkeinspeksjon (trafikkanalyse) Konfigurasjon er deaktivert, vennligst aktiver", + "invalid_export_configuration": "Eksport (trafikkanalyse) Konfigurasjon er deaktivert, vennligst aktiver", + "not_found": "Edgeos url ikke funnet" + }, "step": { "edge_os_additional_settings": { - "title": "Alternativer for EdgeOS.", - "description": "Sett opp dine overvåkede enheter / grensesnitt og sporede enheter, bruk kommaadskilte verdier, for å fjerne eksisterende verdi, bruk avkrysningsruten.", "data": { + "clear-credentials": "Tydelig legitimasjon", "host": "Vert", - "username": "Brukernavn", "password": "Passord", - "clear-credentials": "Fjern legitimasjon" - } + "username": "Brukernavn" + }, + "description": "Sett opp overv\u00e5kede enheter / grensesnitt og sporede enheter, bruk komma -separerte verdier, for \u00e5 fjerne eksisterende verdi, bruk avkrysningsruten.", + "title": "Alternativer for Edgeos." } - }, - "error": { - "invalid_credentials": "Ugyldig legitimasjon til EdgeOS Router", - "auth_general_error": "Generell feil, vennligst prøv igjen", - "not_found": "EdgeOS URL ikke funnet", - "invalid_export_configuration": "Eksport (trafikkanalyse) konfigurasjon er deaktivert, vennligst aktiver", - "invalid_dpi_configuration": "Konfigurasjon av dyp pakkeinspeksjon (trafikkanalyse) er deaktivert, vennligst aktiver", - "empty_device_data": "Kunne ikke hente enhetsdata fra EdgeOS Router", - "already_configured": "EdgeOS-integrasjon allerede konfigurert med navnet", - "incompatible_version": "Inkompatibel versjon (påkrevd minst v2.0)" } } } diff --git a/custom_components/edgeos/translations/pt-BR.json b/custom_components/edgeos/translations/pt-BR.json index 3e596e0..6fa6550 100644 --- a/custom_components/edgeos/translations/pt-BR.json +++ b/custom_components/edgeos/translations/pt-BR.json @@ -1,52 +1,148 @@ { "config": { + "abort": { + "already_configured": "Integra\u00e7\u00e3o de Edgeos ({nome}) j\u00e1 configurado" + }, + "error": { + "already_configured": "Integra\u00e7\u00e3o Edgeos j\u00e1 configurada com o nome", + "auth_general_error": "Falha geral, por favor tente novamente", + "empty_device_data": "N\u00e3o foi poss\u00edvel recuperar os dados do dispositivo do roteador Edgeos", + "incompatible_version": "Vers\u00e3o de firmware n\u00e3o suportada", + "invalid_credentials": "Credenciais inv\u00e1lidas para o roteador Edgeos", + "invalid_dpi_configuration": "A configura\u00e7\u00e3o profunda da inspe\u00e7\u00e3o de pacotes (an\u00e1lise de tr\u00e1fego) est\u00e1 desativada, ative", + "invalid_export_configuration": "A configura\u00e7\u00e3o de exporta\u00e7\u00e3o (an\u00e1lise de tr\u00e1fego) est\u00e1 desativada, ative", + "not_found": "Edgeos url n\u00e3o encontrado" + }, "step": { "user": { - "title": "Configurar Roteador EdgeOS", - "description": "Configure seu roteador EdgeOS, observe que a integração se comunicará usando a conexão HTTPS", "data": { - "host": "Host", - "username": "Nome de usuário", - "password": "Senha" - } + "host": "Hospedar", + "password": "Senha", + "username": "Nome de usu\u00e1rio" + }, + "description": "Configure seu roteador Edgeos, observe que a integra\u00e7\u00e3o se comunicar\u00e1 usando a conex\u00e3o HTTPS", + "title": "Configurar o roteador Edgeos" + } + } + }, + "entity": { + "binary_sensor": { + "firmware": { + "name": "Atualiza\u00e7\u00e3o de firmware" + }, + "interface_connected": { + "name": "Conectada Conectado" } }, - "abort": { - "already_configured": "Integração EdgeOS ({name}) já configurada" + "device_tracker": { + "device_tracker": { + "name": "Status" + } }, - "error": { - "invalid_credentials": "Credenciais inválidas para o roteador EdgeOS", - "auth_general_error": "Falha geral, tente novamente", - "not_found": "URL do EdgeOS não encontrada", - "invalid_export_configuration": "A configuração de exportação (análise de tráfego) está desativada, ative", - "invalid_dpi_configuration": "A configuração de inspeção profunda de pacotes (análise de tráfego) está desativada, ative", - "empty_device_data": "Não foi possível recuperar os dados do dispositivo do roteador EdgeOS", - "already_configured": "EdgeOS integration already configured with the name", - "incompatible_version": "Versão de firmware não suportada" + "number": { + "consider_away_interval": { + "name": "Considere o intervalo fora" + }, + "update_api_interval": { + "name": "Atualizar intervalo da API" + }, + "update_entities_interval": { + "name": "Atualizar intervalos de entidades" + } + }, + "sensor": { + "cpu_usage": { + "name": "Utiliza\u00e7\u00e3o do CPU" + }, + "device_received_rate": { + "name": "Taxa recebeu" + }, + "device_received_traffic": { + "name": "Tr\u00e1fego recebeu" + }, + "device_sent_rate": { + "name": "Taxa enviada" + }, + "device_sent_traffic": { + "name": "Enviou tr\u00e1fego" + }, + "interface_received_dropped": { + "name": "Recebido caiu" + }, + "interface_received_errors": { + "name": "Erros recebidos" + }, + "interface_received_packets": { + "name": "Pacotes recebidos" + }, + "interface_received_rate": { + "name": "Taxa recebeu" + }, + "interface_received_traffic": { + "name": "Tr\u00e1fego recebeu" + }, + "interface_sent_dropped": { + "name": "Enviado caiu" + }, + "interface_sent_errors": { + "name": "Erros enviados" + }, + "interface_sent_packets": { + "name": "Pacotes enviados" + }, + "interface_sent_rate": { + "name": "Taxa enviada" + }, + "interface_sent_traffic": { + "name": "Enviou tr\u00e1fego" + }, + "last_restart": { + "name": "\u00daltima reinicializa\u00e7\u00e3o" + }, + "ram_usage": { + "name": "Uso da RAM" + }, + "unknown_devices": { + "name": "Dispositivos desconhecidos" + } + }, + "switch": { + "device_monitored": { + "name": "Monitorou" + }, + "interface_monitored": { + "name": "Monitorou" + }, + "interface_status": { + "name": "Status" + }, + "log_incoming_messages": { + "name": "Registrar mensagens recebidas" + } } }, "options": { + "error": { + "already_configured": "Integra\u00e7\u00e3o Edgeos j\u00e1 configurada com o nome", + "auth_general_error": "Falha geral, por favor tente novamente", + "empty_device_data": "N\u00e3o foi poss\u00edvel recuperar os dados do dispositivo do roteador Edgeos", + "incompatible_version": "Vers\u00e3o de firmware n\u00e3o suportada", + "invalid_credentials": "Credenciais inv\u00e1lidas para o roteador Edgeos", + "invalid_dpi_configuration": "A configura\u00e7\u00e3o profunda da inspe\u00e7\u00e3o de pacotes (an\u00e1lise de tr\u00e1fego) est\u00e1 desativada, ative", + "invalid_export_configuration": "A configura\u00e7\u00e3o de exporta\u00e7\u00e3o (an\u00e1lise de tr\u00e1fego) est\u00e1 desativada, ative", + "not_found": "Edgeos url n\u00e3o encontrado" + }, "step": { "edge_os_additional_settings": { - "title": "Opções para EdgeOS.", - "description": "Configure seus dispositivos / interfaces monitorados e dispositivos rastreados, use valores separados por vírgulas, para limpar o valor existente, use a caixa de seleção.", "data": { - "host": "Host", - "username": "Nome de usuário", + "clear-credentials": "Credenciais claras", + "host": "Hospedar", "password": "Senha", - "clear-credentials": "Limpar credenciais" - } + "username": "Nome de usu\u00e1rio" + }, + "description": "Configure seus dispositivos / interfaces monitorados e dispositivos rastreados, use valores separados por v\u00edrgula, para limpar o valor existente, use a caixa de sele\u00e7\u00e3o.", + "title": "Op\u00e7\u00f5es para Edgeos." } - }, - "error": { - "invalid_credentials": "Credenciais inválidas para o roteador EdgeOS", - "auth_general_error": "Falha geral, tente novamente", - "not_found": "URL do EdgeOS não encontrado", - "invalid_export_configuration": "A configuração de exportação (análise de tráfego) está desativada, ative", - "invalid_dpi_configuration": "A configuração de inspeção profunda de pacotes (análise de tráfego) está desativada, ative", - "empty_device_data": "Não foi possível recuperar os dados do dispositivo do roteador EdgeOS", - "already_configured": "Integração EdgeOS já configurada com o nome", - "incompatible_version": "Versão de firmware não suportada" } } } diff --git a/requirements.txt b/requirements.txt index 6c2c06c..b37081f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,8 @@ voluptuous aiohttp cryptography numpy + +flatten_json +googletrans==4.0.0rc1 +translators~= 5.4 +deep-translator~=1.9 diff --git a/utils/generate_translations.py b/utils/generate_translations.py new file mode 100644 index 0000000..eec20ef --- /dev/null +++ b/utils/generate_translations.py @@ -0,0 +1,171 @@ +import asyncio +import json +import logging +import os +from pathlib import Path +import sys + +from flatten_json import flatten, unflatten +import translators as ts + +from custom_components.edgeos.common.consts import DOMAIN + +DEBUG = str(os.environ.get("DEBUG", False)).lower() == str(True).lower() + +log_level = logging.DEBUG if DEBUG else logging.INFO + +root = logging.getLogger() +root.setLevel(log_level) + +logging.getLogger("urllib3").setLevel(logging.WARNING) + +stream_handler = logging.StreamHandler(sys.stdout) +stream_handler.setLevel(log_level) +formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s") +stream_handler.setFormatter(formatter) +root.addHandler(stream_handler) + +_LOGGER = logging.getLogger(__name__) + +SOURCE_LANGUAGE = "en" +DESTINATION_LANGUAGES = { + "en": "en", + "nb": "no", + "pt-BR": "pt" +} + +TRANSLATION_PROVIDER = "google" +FLAT_SEPARATOR = "." + + +class TranslationGenerator: + def __init__(self): + self._source_translations = self._get_source_translations() + + self._destinations = DESTINATION_LANGUAGES + + async def initialize(self): + values = flatten(self._source_translations, FLAT_SEPARATOR) + value_keys = list(values.keys()) + last_key = value_keys[len(value_keys) - 1] + + _LOGGER.info( + f"Process will translate {len(values)} sentences " + f"to {len(list(self._destinations.keys()))} languages" + ) + + for lang in self._destinations: + original_values = values.copy() + translated_data = self._get_translations(lang) + translated_values = flatten(translated_data, FLAT_SEPARATOR) + + provider_lang = self._destinations[lang] + lang_cache = {} + + lang_title = provider_lang.upper() + + for key in original_values: + english_value = original_values[key] + + if not isinstance(english_value, str): + continue + + if key in translated_values: + translated_value = translated_values[key] + + _LOGGER.debug( + f"Skip translation to '{lang_title}', " + f"translation of '{english_value}' already exists - '{translated_value}'" + ) + + continue + + if english_value in lang_cache: + translated_value = lang_cache[english_value] + + _LOGGER.debug( + f"Skip translation to '{lang_title}', " + f"translation of '{english_value}' available in cache - {translated_value}" + ) + + elif lang == SOURCE_LANGUAGE: + translated_value = english_value + + _LOGGER.debug( + f"Skip translation to '{lang_title}', " + f"source and destination languages are the same - {translated_value}" + ) + + else: + original_english_value = english_value + + sleep_seconds = 10 if last_key == key else 0 + + translated_value = ts.translate_text( + english_value, + translator=TRANSLATION_PROVIDER, + to_language=provider_lang, + sleep_seconds=sleep_seconds + ) + + lang_cache[english_value] = translated_value + + _LOGGER.debug(f"Translating '{original_english_value}' to {lang_title}: {translated_value}") + + translated_values[key] = translated_value + + translated_data = unflatten(translated_values, FLAT_SEPARATOR) + + self._save_translations(lang, translated_data) + + @staticmethod + def _get_source_translations() -> dict: + current_path = Path(__file__) + parent_directory = current_path.parents[1] + file_path = os.path.join(parent_directory, "custom_components", DOMAIN, "strings.json") + + with open(file_path) as f: + data = json.load(f) + + return data + + @staticmethod + def _get_translations(lang: str): + current_path = Path(__file__) + parent_directory = current_path.parents[1] + file_path = os.path.join(parent_directory, "custom_components", DOMAIN, "translations", f"{lang}.json") + + if os.path.exists(file_path): + with open(file_path) as file: + data = json.load(file) + else: + data = {} + + return data + + @staticmethod + def _save_translations(lang: str, data: dict): + current_path = Path(__file__) + parent_directory = current_path.parents[1] + file_path = os.path.join(parent_directory, "custom_components", DOMAIN, "translations", f"{lang}.json") + + with open(file_path, "w+") as file: + file.write(json.dumps(data, indent=4)) + + _LOGGER.info(f"Translation for {lang.upper()} stored") + + +if __name__ == "__main__": + loop = asyncio.new_event_loop() + + instance = TranslationGenerator() + + try: + loop.create_task(instance.initialize()) + loop.run_forever() + + except KeyboardInterrupt: + _LOGGER.info("Aborted") + + except Exception as rex: + _LOGGER.error(f"Error: {rex}")