core/homeassistant/components/asuswrt/bridge.py

"""aioasuswrt and pyasuswrt bridge classes."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections import namedtuple
from collections.abc import Awaitable, Callable, Coroutine
import functools
import logging
from typing import Any, cast

from aioasuswrt.asuswrt import AsusWrt as AsusWrtLegacy
from aiohttp import ClientSession
from asusrouter import AsusRouter, AsusRouterError
from asusrouter.config import ARConfigKey
from asusrouter.modules.client import AsusClient
from asusrouter.modules.data import AsusData
from asusrouter.modules.homeassistant import convert_to_ha_data, convert_to_ha_sensors
from asusrouter.tools.connection import get_cookie_jar

from homeassistant.const import (
    CONF_HOST,
    CONF_MODE,
    CONF_PASSWORD,
    CONF_PORT,
    CONF_PROTOCOL,
    CONF_USERNAME,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers.aiohttp_client import async_create_clientsession
from homeassistant.helpers.device_registry import format_mac
from homeassistant.helpers.update_coordinator import UpdateFailed

from .const import (
    CONF_DNSMASQ,
    CONF_INTERFACE,
    CONF_REQUIRE_IP,
    CONF_SSH_KEY,
    DEFAULT_DNSMASQ,
    DEFAULT_INTERFACE,
    KEY_METHOD,
    KEY_SENSORS,
    PROTOCOL_HTTP,
    PROTOCOL_HTTPS,
    PROTOCOL_TELNET,
    SENSORS_BYTES,
    SENSORS_LOAD_AVG,
    SENSORS_MEMORY,
    SENSORS_RATES,
    SENSORS_TEMPERATURES_LEGACY,
    SENSORS_UPTIME,
)
from .helpers import clean_dict, translate_to_legacy

SENSORS_TYPE_BYTES = "sensors_bytes"
SENSORS_TYPE_COUNT = "sensors_count"
SENSORS_TYPE_CPU = "sensors_cpu"
SENSORS_TYPE_LOAD_AVG = "sensors_load_avg"
SENSORS_TYPE_MEMORY = "sensors_memory"
SENSORS_TYPE_RATES = "sensors_rates"
SENSORS_TYPE_TEMPERATURES = "sensors_temperatures"
SENSORS_TYPE_UPTIME = "sensors_uptime"

WrtDevice = namedtuple("WrtDevice", ["ip", "name", "connected_to"])  # noqa: PYI024

_LOGGER = logging.getLogger(__name__)

type _FuncType[_T] = Callable[[_T], Awaitable[list[Any] | tuple[Any] | dict[str, Any]]]
type _ReturnFuncType[_T] = Callable[[_T], Coroutine[Any, Any, dict[str, Any]]]


def handle_errors_and_zip[_AsusWrtBridgeT: AsusWrtBridge](
    exceptions: type[Exception] | tuple[type[Exception], ...], keys: list[str] | None
) -> Callable[[_FuncType[_AsusWrtBridgeT]], _ReturnFuncType[_AsusWrtBridgeT]]:
    """Run library methods and zip results or manage exceptions."""

    def _handle_errors_and_zip(
        func: _FuncType[_AsusWrtBridgeT],
    ) -> _ReturnFuncType[_AsusWrtBridgeT]:
        """Run library methods and zip results or manage exceptions."""

        @functools.wraps(func)
        async def _wrapper(self: _AsusWrtBridgeT) -> dict[str, Any]:
            try:
                data = await func(self)
            except exceptions as exc:
                raise UpdateFailed(exc) from exc

            if keys is None:
                if not isinstance(data, dict):
                    raise UpdateFailed("Received invalid data type")
                return data

            if isinstance(data, dict):
                return dict(zip(keys, list(data.values()), strict=False))
            if not isinstance(data, (list, tuple)):
                raise UpdateFailed("Received invalid data type")
            return dict(zip(keys, data, strict=False))

        return _wrapper

    return _handle_errors_and_zip


class AsusWrtBridge(ABC):
    """The Base Bridge abstract class."""

    @staticmethod
    def get_bridge(
        hass: HomeAssistant, conf: dict[str, Any], options: dict[str, Any] | None = None
    ) -> AsusWrtBridge:
        """Get Bridge instance."""
        if conf[CONF_PROTOCOL] in (PROTOCOL_HTTPS, PROTOCOL_HTTP):
            session = async_create_clientsession(
                hass,
                cookie_jar=get_cookie_jar(),
            )
            return AsusWrtHttpBridge(conf, session)
        return AsusWrtLegacyBridge(conf, options)

    def __init__(self, host: str) -> None:
        """Initialize Bridge."""
        self._configuration_url = f"http://{host}"
        self._host = host
        self._firmware: str | None = None
        self._label_mac: str | None = None
        self._model: str | None = None
        self._model_id: str | None = None
        self._serial_number: str | None = None

    @property
    def configuration_url(self) -> str:
        """Return configuration URL."""
        return self._configuration_url

    @property
    def host(self) -> str:
        """Return hostname."""
        return self._host

    @property
    def firmware(self) -> str | None:
        """Return firmware information."""
        return self._firmware

    @property
    def label_mac(self) -> str | None:
        """Return label mac information."""
        return self._label_mac

    @property
    def model(self) -> str | None:
        """Return model information."""
        return self._model

    @property
    def model_id(self) -> str | None:
        """Return model_id information."""
        return self._model_id

    @property
    def serial_number(self) -> str | None:
        """Return serial number information."""
        return self._serial_number

    @property
    @abstractmethod
    def is_connected(self) -> bool:
        """Get connected status."""

    @abstractmethod
    async def async_connect(self) -> None:
        """Connect to the device."""

    @abstractmethod
    async def async_disconnect(self) -> None:
        """Disconnect to the device."""

    @abstractmethod
    async def async_get_connected_devices(self) -> dict[str, WrtDevice]:
        """Get list of connected devices."""

    @abstractmethod
    async def async_get_available_sensors(self) -> dict[str, dict[str, Any]]:
        """Return a dictionary of available sensors for this bridge."""


class AsusWrtLegacyBridge(AsusWrtBridge):
    """The Bridge that use legacy library."""

    def __init__(
        self, conf: dict[str, Any], options: dict[str, Any] | None = None
    ) -> None:
        """Initialize Bridge."""
        super().__init__(conf[CONF_HOST])
        self._protocol: str = conf[CONF_PROTOCOL]
        self._api: AsusWrtLegacy = self._get_api(conf, options)

    @staticmethod
    def _get_api(
        conf: dict[str, Any], options: dict[str, Any] | None = None
    ) -> AsusWrtLegacy:
        """Get the AsusWrtLegacy API."""
        opt = options or {}

        return AsusWrtLegacy(
            conf[CONF_HOST],
            conf.get(CONF_PORT),
            conf[CONF_PROTOCOL] == PROTOCOL_TELNET,
            conf[CONF_USERNAME],
            conf.get(CONF_PASSWORD, ""),
            conf.get(CONF_SSH_KEY, ""),
            conf[CONF_MODE],
            opt.get(CONF_REQUIRE_IP, True),
            interface=opt.get(CONF_INTERFACE, DEFAULT_INTERFACE),
            dnsmasq=opt.get(CONF_DNSMASQ, DEFAULT_DNSMASQ),
        )

    @property
    def is_connected(self) -> bool:
        """Get connected status."""
        return cast(bool, self._api.is_connected)

    async def async_connect(self) -> None:
        """Connect to the device."""
        await self._api.connection.async_connect()

        # get main router properties
        if self._label_mac is None:
            await self._get_label_mac()
        if self._firmware is None:
            await self._get_firmware()
        if self._model is None:
            await self._get_model()

    async def async_disconnect(self) -> None:
        """Disconnect to the device."""
        if self._api is not None and self._protocol == PROTOCOL_TELNET:
            self._api.connection.disconnect()

    async def async_get_connected_devices(self) -> dict[str, WrtDevice]:
        """Get list of connected devices."""
        api_devices = await self._api.async_get_connected_devices()
        return {
            format_mac(mac): WrtDevice(dev.ip, dev.name, None)
            for mac, dev in api_devices.items()
        }

    async def _get_nvram_info(self, info_type: str) -> dict[str, Any]:
        """Get AsusWrt router info from nvram."""
        info = {}
        try:
            info = await self._api.async_get_nvram(info_type)
        except OSError as exc:
            _LOGGER.warning(
                "Error calling method async_get_nvram(%s): %s", info_type, exc
            )

        return info

    async def _get_label_mac(self) -> None:
        """Get label mac information."""
        label_mac = await self._get_nvram_info("LABEL_MAC")
        if label_mac and "label_mac" in label_mac:
            self._label_mac = format_mac(label_mac["label_mac"])

    async def _get_firmware(self) -> None:
        """Get firmware information."""
        firmware = await self._get_nvram_info("FIRMWARE")
        if firmware and "firmver" in firmware:
            firmver: str = firmware["firmver"]
            if "buildno" in firmware:
                firmver += f" (build {firmware['buildno']})"
            self._firmware = firmver

    async def _get_model(self) -> None:
        """Get model information."""
        model = await self._get_nvram_info("MODEL")
        if model and "model" in model:
            self._model = model["model"]

    async def async_get_available_sensors(self) -> dict[str, dict[str, Any]]:
        """Return a dictionary of available sensors for this bridge."""
        sensors_temperatures = await self._get_available_temperature_sensors()
        return {
            SENSORS_TYPE_BYTES: {
                KEY_SENSORS: SENSORS_BYTES,
                KEY_METHOD: self._get_bytes,
            },
            SENSORS_TYPE_LOAD_AVG: {
                KEY_SENSORS: SENSORS_LOAD_AVG,
                KEY_METHOD: self._get_load_avg,
            },
            SENSORS_TYPE_RATES: {
                KEY_SENSORS: SENSORS_RATES,
                KEY_METHOD: self._get_rates,
            },
            SENSORS_TYPE_TEMPERATURES: {
                KEY_SENSORS: sensors_temperatures,
                KEY_METHOD: self._get_temperatures,
            },
        }

    async def _get_available_temperature_sensors(self) -> list[str]:
        """Check which temperature information is available on the router."""
        availability = await self._api.async_find_temperature_commands()
        return [SENSORS_TEMPERATURES_LEGACY[i] for i in range(3) if availability[i]]

    @handle_errors_and_zip((IndexError, OSError, ValueError), SENSORS_BYTES)
    async def _get_bytes(self) -> Any:
        """Fetch byte information from the router."""
        return await self._api.async_get_bytes_total()

    @handle_errors_and_zip((IndexError, OSError, ValueError), SENSORS_RATES)
    async def _get_rates(self) -> Any:
        """Fetch rates information from the router."""
        return await self._api.async_get_current_transfer_rates()

    @handle_errors_and_zip((IndexError, OSError, ValueError), SENSORS_LOAD_AVG)
    async def _get_load_avg(self) -> Any:
        """Fetch load average information from the router."""
        return await self._api.async_get_loadavg()

    @handle_errors_and_zip((OSError, ValueError), None)
    async def _get_temperatures(self) -> Any:
        """Fetch temperatures information from the router."""
        return await self._api.async_get_temperature()


class AsusWrtHttpBridge(AsusWrtBridge):
    """The Bridge that use HTTP library."""

    def __init__(self, conf: dict[str, Any], session: ClientSession) -> None:
        """Initialize Bridge that use HTTP library."""
        super().__init__(conf[CONF_HOST])
        # Get API configuration
        config = self._get_api_config()
        self._api = self._get_api(conf, session, config)

    @staticmethod
    def _get_api(
        conf: dict[str, Any], session: ClientSession, config: dict[ARConfigKey, Any]
    ) -> AsusRouter:
        """Get the AsusRouter API."""
        return AsusRouter(
            hostname=conf[CONF_HOST],
            username=conf[CONF_USERNAME],
            password=conf.get(CONF_PASSWORD, ""),
            use_ssl=conf[CONF_PROTOCOL] == PROTOCOL_HTTPS,
            port=conf.get(CONF_PORT),
            session=session,
            config=config,
        )

    def _get_api_config(self) -> dict[ARConfigKey, Any]:
        """Get configuration for the API."""
        return {
            # Enable automatic temperature data correction in the library
            ARConfigKey.OPTIMISTIC_TEMPERATURE: True,
            # Disable `warning`-level log message when temperature
            # is corrected by setting it to already notified.
            ARConfigKey.NOTIFIED_OPTIMISTIC_TEMPERATURE: True,
        }

    @property
    def is_connected(self) -> bool:
        """Get connected status."""
        return self._api.connected

    async def async_connect(self) -> None:
        """Connect to the device."""
        await self._api.async_connect()

        # Collect the identity
        _identity = await self._api.async_get_identity()

        # get main router properties
        if mac := _identity.mac:
            self._label_mac = format_mac(mac)
        self._configuration_url = self._api.webpanel
        self._firmware = str(_identity.firmware)
        self._model = _identity.model
        self._model_id = _identity.product_id
        self._serial_number = _identity.serial

    async def async_disconnect(self) -> None:
        """Disconnect to the device."""
        await self._api.async_disconnect()

    async def _get_data(
        self,
        datatype: AsusData,
        force: bool = False,
    ) -> dict[str, Any]:
        """Get data from the device.

        This is a generic method which automatically converts to
        the Home Assistant-compatible format.
        """
        try:
            raw = await self._api.async_get_data(datatype, force=force)
            return translate_to_legacy(clean_dict(convert_to_ha_data(raw)))
        except AsusRouterError as ex:
            raise UpdateFailed(ex) from ex

    async def _get_sensors(self, datatype: AsusData) -> list[str]:
        """Get the available sensors.

        This is a generic method which automatically converts to
        the Home Assistant-compatible format.
        """
        sensors = []
        try:
            data = await self._api.async_get_data(datatype)
            # Get the list of sensors from the raw data
            # and translate in to the legacy format
            sensors = translate_to_legacy(convert_to_ha_sensors(data, datatype))
            _LOGGER.debug("Available `%s` sensors: %s", datatype.value, sensors)
        except AsusRouterError as ex:
            _LOGGER.warning(
                "Cannot get available `%s` sensors with exception: %s",
                datatype.value,
                ex,
            )
        return sensors

    async def async_get_connected_devices(self) -> dict[str, WrtDevice]:
        """Get list of connected devices."""
        api_devices: dict[str, AsusClient] = await self._api.async_get_data(
            AsusData.CLIENTS, force=True
        )
        return {
            format_mac(mac): WrtDevice(
                dev.connection.ip_address, dev.description.name, dev.connection.node
            )
            for mac, dev in api_devices.items()
            if dev.connection is not None
            and dev.description is not None
            and dev.connection.ip_address is not None
        }

    async def async_get_available_sensors(self) -> dict[str, dict[str, Any]]:
        """Return a dictionary of available sensors for this bridge."""
        return {
            SENSORS_TYPE_BYTES: {
                KEY_SENSORS: SENSORS_BYTES,
                KEY_METHOD: self._get_bytes,
            },
            SENSORS_TYPE_CPU: {
                KEY_SENSORS: await self._get_sensors(AsusData.CPU),
                KEY_METHOD: self._get_cpu_usage,
            },
            SENSORS_TYPE_LOAD_AVG: {
                KEY_SENSORS: await self._get_sensors(AsusData.SYSINFO),
                KEY_METHOD: self._get_load_avg,
            },
            SENSORS_TYPE_MEMORY: {
                KEY_SENSORS: SENSORS_MEMORY,
                KEY_METHOD: self._get_memory_usage,
            },
            SENSORS_TYPE_RATES: {
                KEY_SENSORS: SENSORS_RATES,
                KEY_METHOD: self._get_rates,
            },
            SENSORS_TYPE_UPTIME: {
                KEY_SENSORS: SENSORS_UPTIME,
                KEY_METHOD: self._get_uptime,
            },
            SENSORS_TYPE_TEMPERATURES: {
                KEY_SENSORS: await self._get_sensors(AsusData.TEMPERATURE),
                KEY_METHOD: self._get_temperatures,
            },
        }

    async def _get_bytes(self) -> Any:
        """Fetch byte information from the router."""
        return await self._get_data(AsusData.NETWORK)

    async def _get_rates(self) -> Any:
        """Fetch rates information from the router."""
        data = await self._get_data(AsusData.NETWORK)
        # Convert from bits/s to Bytes/s for compatibility with legacy sensors
        return {
            key: (
                value / 8
                if key in SENSORS_RATES and isinstance(value, (int, float))
                else value
            )
            for key, value in data.items()
        }

    async def _get_load_avg(self) -> Any:
        """Fetch cpu load avg information from the router."""
        return await self._get_data(AsusData.SYSINFO)

    async def _get_temperatures(self) -> Any:
        """Fetch temperatures information from the router."""
        return await self._get_data(AsusData.TEMPERATURE)

    async def _get_cpu_usage(self) -> Any:
        """Fetch cpu information from the router."""
        return await self._get_data(AsusData.CPU)

    async def _get_memory_usage(self) -> Any:
        """Fetch memory information from the router."""
        return await self._get_data(AsusData.RAM)

    async def _get_uptime(self) -> dict[str, Any]:
        """Fetch uptime from the router."""
        return await self._get_data(AsusData.BOOTTIME)