core/homeassistant/components/zha/radio_manager.py

"""ZHA radio manager."""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
import contextlib
from contextlib import suppress
import copy
import enum
import logging
import os
from typing import Any, Self

from bellows.config import CONF_USE_THREAD
import voluptuous as vol
from zha.application.const import RadioType
from zigpy.application import ControllerApplication
import zigpy.backups
from zigpy.config import (
    CONF_DATABASE,
    CONF_DEVICE,
    CONF_DEVICE_PATH,
    CONF_NWK_BACKUP_ENABLED,
    SCHEMA_DEVICE,
)
from zigpy.exceptions import NetworkNotFormed

from homeassistant import config_entries
from homeassistant.components import usb
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.service_info.usb import UsbServiceInfo

from . import repairs
from .const import (
    CONF_RADIO_TYPE,
    CONF_ZIGPY,
    DEFAULT_DATABASE_NAME,
    EZSP_OVERWRITE_EUI64,
)
from .helpers import get_zha_data

RECOMMENDED_RADIOS = (
    RadioType.ezsp,
    RadioType.znp,
    RadioType.deconz,
)

# Only the common radio types will be autoprobed, ordered by new device popularity.
# XBee takes too long to probe since it scans through all possible bauds and likely has
# very few users to begin with.
AUTOPROBE_RADIOS = RECOMMENDED_RADIOS

CONNECT_DELAY_S = 1.0
RETRY_DELAY_S = 1.0

BACKUP_RETRIES = 5
MIGRATION_RETRIES = 100


DEVICE_SCHEMA = vol.Schema(
    {
        vol.Required("path"): str,
        vol.Optional("baudrate", default=115200): int,
        vol.Optional("flow_control", default=None): vol.In(
            ["hardware", "software", None]
        ),
    }
)

HARDWARE_DISCOVERY_SCHEMA = vol.Schema(
    {
        vol.Required("name"): str,
        vol.Required("port"): DEVICE_SCHEMA,
        vol.Required("radio_type"): str,
    }
)

HARDWARE_MIGRATION_SCHEMA = vol.Schema(
    {
        vol.Required("new_discovery_info"): HARDWARE_DISCOVERY_SCHEMA,
        vol.Required("old_discovery_info"): vol.Schema(
            {
                vol.Exclusive("hw", "discovery"): HARDWARE_DISCOVERY_SCHEMA,
                vol.Exclusive("usb", "discovery"): UsbServiceInfo,
            }
        ),
    }
)

_LOGGER = logging.getLogger(__name__)


class ProbeResult(enum.StrEnum):
    """Radio firmware probing result."""

    RADIO_TYPE_DETECTED = "radio_type_detected"
    WRONG_FIRMWARE_INSTALLED = "wrong_firmware_installed"
    PROBING_FAILED = "probing_failed"


def _allow_overwrite_ezsp_ieee(
    backup: zigpy.backups.NetworkBackup,
) -> zigpy.backups.NetworkBackup:
    """Return a new backup with the flag to allow overwriting the EZSP EUI64."""
    new_stack_specific = copy.deepcopy(backup.network_info.stack_specific)
    new_stack_specific.setdefault("ezsp", {})[EZSP_OVERWRITE_EUI64] = True

    return backup.replace(
        network_info=backup.network_info.replace(stack_specific=new_stack_specific)
    )


def _prevent_overwrite_ezsp_ieee(
    backup: zigpy.backups.NetworkBackup,
) -> zigpy.backups.NetworkBackup:
    """Return a new backup without the flag to allow overwriting the EZSP EUI64."""
    if "ezsp" not in backup.network_info.stack_specific:
        return backup

    new_stack_specific = copy.deepcopy(backup.network_info.stack_specific)
    new_stack_specific.setdefault("ezsp", {}).pop(EZSP_OVERWRITE_EUI64, None)

    return backup.replace(
        network_info=backup.network_info.replace(stack_specific=new_stack_specific)
    )


class ZhaRadioManager:
    """Helper class with radio related functionality."""

    hass: HomeAssistant

    def __init__(self) -> None:
        """Initialize ZhaRadioManager instance."""
        self.device_path: str | None = None
        self.device_settings: dict[str, Any] | None = None
        self.radio_type: RadioType | None = None
        self.current_settings: zigpy.backups.NetworkBackup | None = None
        self.backups: list[zigpy.backups.NetworkBackup] = []
        self.chosen_backup: zigpy.backups.NetworkBackup | None = None

    @classmethod
    def from_config_entry(
        cls, hass: HomeAssistant, config_entry: config_entries.ConfigEntry
    ) -> Self:
        """Create an instance from a config entry."""
        mgr = cls()
        mgr.hass = hass
        mgr.device_path = config_entry.data[CONF_DEVICE][CONF_DEVICE_PATH]
        mgr.device_settings = config_entry.data[CONF_DEVICE]
        mgr.radio_type = RadioType[config_entry.data[CONF_RADIO_TYPE]]

        return mgr

    @property
    def zigpy_database_path(self) -> str:
        """Path to `zigbee.db`."""
        config = get_zha_data(self.hass).yaml_config

        return config.get(
            CONF_DATABASE,
            self.hass.config.path(DEFAULT_DATABASE_NAME),
        )

    @contextlib.asynccontextmanager
    async def create_zigpy_app(
        self, *, connect: bool = True
    ) -> AsyncIterator[ControllerApplication]:
        """Connect to the radio with the current config and then clean up."""
        assert self.radio_type is not None

        config = get_zha_data(self.hass).yaml_config
        app_config = config.get(CONF_ZIGPY, {}).copy()

        database_path: str | None = self.zigpy_database_path

        # Don't create `zigbee.db` if it doesn't already exist
        try:
            if database_path is not None and not await self.hass.async_add_executor_job(
                os.path.exists, database_path
            ):
                database_path = None
        except OSError as error:
            raise HomeAssistantError(
                f"Could not read the ZHA database {database_path}: {error}"
            ) from error

        app_config[CONF_DATABASE] = database_path
        app_config[CONF_DEVICE] = self.device_settings
        app_config[CONF_NWK_BACKUP_ENABLED] = False
        app_config[CONF_USE_THREAD] = False

        app = await self.radio_type.controller.new(
            app_config, auto_form=False, start_radio=False
        )

        try:
            if connect:
                try:
                    await app.connect()
                except OSError as error:
                    raise HomeAssistantError(
                        f"Failed to connect to Zigbee adapter: {error}"
                    ) from error

            yield app
        finally:
            await app.shutdown()
            await asyncio.sleep(CONNECT_DELAY_S)

    async def restore_backup(
        self,
        backup: zigpy.backups.NetworkBackup | None = None,
        *,
        overwrite_ieee: bool = False,
        **kwargs: Any,
    ) -> None:
        """Restore the provided network backup, passing through kwargs."""
        if backup is None:
            backup = self.chosen_backup

        assert backup is not None

        if self.current_settings is not None and self.current_settings.supersedes(
            backup
        ):
            return

        if overwrite_ieee:
            backup = _allow_overwrite_ezsp_ieee(backup)

        async with self.create_zigpy_app() as app:
            await app.can_write_network_settings(
                network_info=backup.network_info,
                node_info=backup.node_info,
            )
            await app.backups.restore_backup(backup, **kwargs)

    @staticmethod
    def parse_radio_type(radio_type: str) -> RadioType:
        """Parse a radio type name, accounting for past aliases."""
        if radio_type == "efr32":
            return RadioType.ezsp

        return RadioType[radio_type]

    async def detect_radio_type(self) -> ProbeResult:
        """Probe all radio types on the current port."""
        assert self.device_path is not None

        for radio in AUTOPROBE_RADIOS:
            _LOGGER.debug("Attempting to probe radio type %s", radio)

            dev_config = SCHEMA_DEVICE({CONF_DEVICE_PATH: self.device_path})
            probe_result = await radio.controller.probe(dev_config)

            if not probe_result:
                continue

            # Radio library probing can succeed and return new device settings
            if isinstance(probe_result, dict):
                dev_config = probe_result

            self.radio_type = radio
            self.device_settings = dev_config

            repairs.async_delete_blocking_issues(self.hass)
            return ProbeResult.RADIO_TYPE_DETECTED

        with suppress(repairs.wrong_silabs_firmware.AlreadyRunningEZSP):
            if await repairs.wrong_silabs_firmware.warn_on_wrong_silabs_firmware(
                self.hass, self.device_path
            ):
                return ProbeResult.WRONG_FIRMWARE_INSTALLED

        return ProbeResult.PROBING_FAILED

    async def _async_read_backups_from_database(
        self,
    ) -> list[zigpy.backups.NetworkBackup]:
        """Read the list of backups from the database, internal."""
        async with self.create_zigpy_app(connect=False) as app:
            backups = app.backups.backups.copy()
            backups.sort(reverse=True, key=lambda b: b.backup_time)

        return backups

    async def async_read_backups_from_database(self) -> None:
        """Read the list of backups from the database."""
        self.backups = await self._async_read_backups_from_database()

    async def async_load_network_settings(
        self, *, create_backup: bool = False
    ) -> zigpy.backups.NetworkBackup | None:
        """Connect to the radio and load its current network settings."""
        backup = None

        async with self.create_zigpy_app() as app:
            # Check if the stick has any settings and load them
            try:
                await app.load_network_info()
            except NetworkNotFormed:
                pass
            else:
                self.current_settings = zigpy.backups.NetworkBackup(
                    network_info=app.state.network_info,
                    node_info=app.state.node_info,
                )

                if create_backup:
                    backup = await app.backups.create_backup()

            # The list of backups will always exist
            self.backups = app.backups.backups.copy()
            self.backups.sort(reverse=True, key=lambda b: b.backup_time)

        return backup

    async def async_form_network(self) -> None:
        """Form a brand-new network."""

        # When forming a new network, we delete the ZHA database to prevent old devices
        # from appearing in an unusable state
        with suppress(OSError):
            await self.hass.async_add_executor_job(os.remove, self.zigpy_database_path)

        async with self.create_zigpy_app() as app:
            await app.form_network()

    async def async_reset_adapter(self) -> None:
        """Reset the current adapter."""
        async with self.create_zigpy_app() as app:
            await app.reset_network_info()


class ZhaMultiPANMigrationHelper:
    """Helper class for automatic migration when upgrading the firmware of a radio.

    This class is currently only intended to be used when changing the firmware on the
    radio used in the Home Assistant SkyConnect USB stick and the Home Assistant Yellow
    from Zigbee only firmware to firmware supporting both Zigbee and Thread.
    """

    def __init__(
        self, hass: HomeAssistant, config_entry: config_entries.ConfigEntry
    ) -> None:
        """Initialize MigrationHelper instance."""
        self._config_entry = config_entry
        self._hass = hass
        self._radio_mgr = ZhaRadioManager()
        self._radio_mgr.hass = hass

    async def async_initiate_migration(self, data: dict[str, Any]) -> bool:
        """Initiate ZHA migration.

        The passed data should contain:
        - Discovery data identifying the device being firmware updated
        - Discovery data for connecting to the device after the firmware update is
          completed.

        Returns True if async_finish_migration should be called after the firmware
        update is completed.
        """
        migration_data = HARDWARE_MIGRATION_SCHEMA(data)

        name = migration_data["new_discovery_info"]["name"]
        new_radio_type = ZhaRadioManager.parse_radio_type(
            migration_data["new_discovery_info"]["radio_type"]
        )

        new_device_settings = SCHEMA_DEVICE(
            migration_data["new_discovery_info"]["port"]
        )

        if "hw" in migration_data["old_discovery_info"]:
            old_device_path = migration_data["old_discovery_info"]["hw"]["port"]["path"]
        else:  # usb
            device = migration_data["old_discovery_info"]["usb"].device
            old_device_path = await self._hass.async_add_executor_job(
                usb.get_serial_by_id, device
            )

        if self._config_entry.data[CONF_DEVICE][CONF_DEVICE_PATH] != old_device_path:
            # ZHA is using another radio, do nothing
            return False

        # OperationNotAllowed: ZHA is not running
        with suppress(config_entries.OperationNotAllowed):
            await self._hass.config_entries.async_unload(self._config_entry.entry_id)

        # Temporarily connect to the old radio to read its settings
        config_entry_data = self._config_entry.data
        old_radio_mgr = ZhaRadioManager()
        old_radio_mgr.hass = self._hass
        old_radio_mgr.device_path = config_entry_data[CONF_DEVICE][CONF_DEVICE_PATH]
        old_radio_mgr.device_settings = config_entry_data[CONF_DEVICE]
        old_radio_mgr.radio_type = RadioType[config_entry_data[CONF_RADIO_TYPE]]

        for retry in range(BACKUP_RETRIES):
            try:
                backup = await old_radio_mgr.async_load_network_settings(
                    create_backup=True
                )
                break
            except OSError as err:
                if retry >= BACKUP_RETRIES - 1:
                    raise

                _LOGGER.debug(
                    "Failed to create backup %r, retrying in %s seconds",
                    err,
                    RETRY_DELAY_S,
                )

            await asyncio.sleep(RETRY_DELAY_S)

        # Then configure the radio manager for the new radio to use the new settings
        self._radio_mgr.chosen_backup = backup
        self._radio_mgr.radio_type = new_radio_type
        self._radio_mgr.device_path = new_device_settings[CONF_DEVICE_PATH]
        self._radio_mgr.device_settings = new_device_settings
        device_settings = self._radio_mgr.device_settings.copy()

        # Update the config entry settings
        self._hass.config_entries.async_update_entry(
            entry=self._config_entry,
            data={
                CONF_DEVICE: device_settings,
                CONF_RADIO_TYPE: self._radio_mgr.radio_type.name,
            },
            options=self._config_entry.options,
            title=name,
        )
        return True

    async def async_finish_migration(self) -> None:
        """Finish ZHA migration.

        Throws an exception if the migration did not succeed.
        """
        # Restore the backup, permanently overwriting the device IEEE address
        for retry in range(MIGRATION_RETRIES):
            try:
                await self._radio_mgr.restore_backup(overwrite_ieee=True)
                break
            except OSError as err:
                if retry >= MIGRATION_RETRIES - 1:
                    raise

                _LOGGER.debug(
                    "Failed to restore backup %r, retrying in %s seconds",
                    err,
                    RETRY_DELAY_S,
                )

            await asyncio.sleep(RETRY_DELAY_S)

        _LOGGER.debug("Restored backup after %s retries", retry)

        # Launch ZHA again
        # OperationNotAllowed: ZHA is not unloaded
        with suppress(config_entries.OperationNotAllowed):
            await self._hass.config_entries.async_setup(self._config_entry.entry_id)