core/homeassistant/components/airos/__init__.py

"""The Ubiquiti airOS integration."""

import logging

from aiohttp import ClientSession, TCPConnector
from airos.airos6 import AirOS6
from airos.airos8 import AirOS8
from airos.exceptions import (
    AirOSConnectionAuthenticationError,
    AirOSConnectionSetupError,
    AirOSDataMissingError,
    AirOSDeviceConnectionError,
    AirOSKeyDataMissingError,
    AirOSTLSCompatibilityError,
)
from airos.helpers import DetectDeviceData, async_get_firmware_data

from homeassistant.const import (
    CONF_HOST,
    CONF_PASSWORD,
    CONF_SSL,
    CONF_USERNAME,
    CONF_VERIFY_SSL,
    Platform,
)
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import (
    ConfigEntryAuthFailed,
    ConfigEntryError,
    ConfigEntryNotReady,
)
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.aiohttp_client import async_get_clientsession

from .const import (
    CONF_LEGACY_SSL,
    DEFAULT_SSL,
    DEFAULT_VERIFY_SSL,
    DOMAIN,
    SECTION_ADDITIONAL_SETTINGS,
)
from .coordinator import (
    AirOSConfigEntry,
    AirOSDataUpdateCoordinator,
    AirOSFirmwareUpdateCoordinator,
    AirOSRuntimeData,
)
from .helpers import build_legacy_context

_PLATFORMS: list[Platform] = [
    Platform.BINARY_SENSOR,
    Platform.BUTTON,
    Platform.SENSOR,
    Platform.UPDATE,
]


_LOGGER = logging.getLogger(__name__)


async def async_setup_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> bool:
    """Set up Ubiquiti airOS from a config entry."""
    owns_session = False
    verify_ssl = entry.data[SECTION_ADDITIONAL_SETTINGS][CONF_VERIFY_SSL]

    # By default airOS 8 comes with self-signed SSL certificates,
    # with no option in the web UI to change or upload a custom certificate.
    session = async_get_clientsession(hass, verify_ssl=verify_ssl)

    if entry.data.get(CONF_LEGACY_SSL, False):
        session = ClientSession(
            connector=TCPConnector(ssl=build_legacy_context(verify_ssl=verify_ssl))
        )
        owns_session = True

    conn_data = {
        CONF_HOST: entry.data[CONF_HOST],
        CONF_USERNAME: entry.data[CONF_USERNAME],
        CONF_PASSWORD: entry.data[CONF_PASSWORD],
        "session": session,
        "use_ssl": entry.data[SECTION_ADDITIONAL_SETTINGS][CONF_SSL],
    }

    async def close_session() -> None:
        """Close legacy session before raising if needed."""
        if owns_session:
            await session.close()

    # Determine firmware version before creating the device instance
    try:
        device_data: DetectDeviceData = await async_get_firmware_data(**conn_data)

    except (
        AirOSConnectionSetupError,
        AirOSDeviceConnectionError,
        AirOSTLSCompatibilityError,
        TimeoutError,
    ) as err:
        await close_session()
        raise ConfigEntryNotReady from err
    except (
        AirOSConnectionAuthenticationError,
        AirOSDataMissingError,
    ) as err:
        await close_session()
        raise ConfigEntryAuthFailed from err
    except AirOSKeyDataMissingError as err:
        await close_session()
        raise ConfigEntryError(
            translation_domain=DOMAIN, translation_key="key_data_missing"
        ) from err
    except Exception as err:
        await close_session()
        raise ConfigEntryError(
            translation_domain=DOMAIN, translation_key="unknown"
        ) from err

    airos_class: type[AirOS8 | AirOS6] = (
        AirOS8 if device_data["fw_major"] == 8 else AirOS6
    )

    airos_device = airos_class(**conn_data)

    data_coordinator = AirOSDataUpdateCoordinator(
        hass, entry, device_data, airos_device
    )

    try:
        await data_coordinator.async_config_entry_first_refresh()

        firmware_coordinator: AirOSFirmwareUpdateCoordinator | None = None
        if device_data["fw_major"] >= 8:
            firmware_coordinator = AirOSFirmwareUpdateCoordinator(
                hass, entry, airos_device
            )
            await firmware_coordinator.async_config_entry_first_refresh()
    except ConfigEntryNotReady, ConfigEntryAuthFailed:
        await close_session()
        raise
    except Exception as err:
        await close_session()
        raise ConfigEntryError(
            translation_domain=DOMAIN, translation_key="unknown"
        ) from err

    entry.runtime_data = AirOSRuntimeData(
        status=data_coordinator,
        firmware=firmware_coordinator,
        owns_session=owns_session,
        session=session,
    )

    await hass.config_entries.async_forward_entry_setups(entry, _PLATFORMS)

    return True


async def async_migrate_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> bool:
    """Migrate old config entry."""

    # 1.1 Migrate config_entry to add additional ssl settings
    if entry.version == 1 and entry.minor_version == 1:
        new_minor_version = 2
        new_data = {**entry.data}
        additional_data = {
            CONF_SSL: DEFAULT_SSL,
            CONF_VERIFY_SSL: DEFAULT_VERIFY_SSL,
        }
        new_data[SECTION_ADDITIONAL_SETTINGS] = additional_data

        hass.config_entries.async_update_entry(
            entry,
            data=new_data,
            minor_version=new_minor_version,
        )

    # 2.1 Migrate binary_sensor entity unique_id from device_id to mac_address
    #     Step 1 - migrate binary_sensor entity unique_id
    #     Step 2 - migrate device entity identifier
    if entry.version == 1:
        new_version = 2
        new_minor_version = 1

        mac_adress = dr.format_mac(entry.unique_id)

        device_registry = dr.async_get(hass)
        if device_entry := device_registry.async_get_device_by_connection(
            (dr.CONNECTION_NETWORK_MAC, mac_adress), entry.entry_id
        ):
            old_device_id = next(
                (
                    device_id
                    for domain, device_id in device_entry.identifiers
                    if domain == DOMAIN
                ),
            )

            @callback
            def update_unique_id(
                entity_entry: er.RegistryEntry,
            ) -> dict[str, str] | None:
                """Update unique id from device_id to mac address."""
                if old_device_id and entity_entry.unique_id.startswith(old_device_id):
                    suffix = entity_entry.unique_id.removeprefix(old_device_id)
                    new_unique_id = f"{mac_adress}{suffix}"
                    return {"new_unique_id": new_unique_id}
                return None

            await er.async_migrate_entries(hass, entry.entry_id, update_unique_id)

            new_identifiers = device_entry.identifiers.copy()
            new_identifiers.discard((DOMAIN, old_device_id))
            new_identifiers.add((DOMAIN, mac_adress))
            device_registry.async_update_device(
                device_entry.id, new_identifiers=new_identifiers
            )

        hass.config_entries.async_update_entry(
            entry, version=new_version, minor_version=new_minor_version
        )

    if entry.version == 2:
        new_version = 3
        new_minor_version = 1
        new_data = {**entry.data}

        if "advanced_settings" in new_data:
            new_data[SECTION_ADDITIONAL_SETTINGS] = new_data.pop("advanced_settings")

        hass.config_entries.async_update_entry(
            entry,
            data=new_data,
            version=new_version,
            minor_version=new_minor_version,
        )

    return True


async def async_unload_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> bool:
    """Unload a config entry."""
    unload_state = await hass.config_entries.async_unload_platforms(entry, _PLATFORMS)
    # Clean up legacy session if needed
    if unload_state and entry.runtime_data.owns_session:
        await entry.runtime_data.session.close()

    return unload_state