"""The NextDNS component.""" import asyncio from dataclasses import dataclass from types import MappingProxyType from aiohttp.client_exceptions import ClientConnectorError from nextdns import ( AnalyticsDnssec, AnalyticsEncryption, AnalyticsIpVersions, AnalyticsProtocols, AnalyticsStatus, ApiError, ConnectionStatus, InvalidApiKeyError, NextDns, Settings, ) from tenacity import RetryError from homeassistant.config_entries import ConfigEntry, ConfigSubentry from homeassistant.const import CONF_API_KEY, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import ( config_validation as cv, device_registry as dr, entity_registry as er, ) from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.typing import ConfigType from .const import ( ATTR_CONNECTION, ATTR_DNSSEC, ATTR_ENCRYPTION, ATTR_IP_VERSIONS, ATTR_PROTOCOLS, ATTR_SETTINGS, ATTR_STATUS, CONF_PROFILE_ID, DOMAIN, SUBENTRY_TYPE_PROFILE, ) from .coordinator import ( NextDnsConnectionUpdateCoordinator, NextDnsDnssecUpdateCoordinator, NextDnsEncryptionUpdateCoordinator, NextDnsIpVersionsUpdateCoordinator, NextDnsProtocolsUpdateCoordinator, NextDnsSettingsUpdateCoordinator, NextDnsStatusUpdateCoordinator, NextDnsUpdateCoordinator, ) type NextDnsConfigEntry = ConfigEntry[NextDnsData] @dataclass class NextDnsCoordinators: """Coordinators for a NextDNS profile.""" connection: NextDnsUpdateCoordinator[ConnectionStatus] dnssec: NextDnsUpdateCoordinator[AnalyticsDnssec] encryption: NextDnsUpdateCoordinator[AnalyticsEncryption] ip_versions: NextDnsUpdateCoordinator[AnalyticsIpVersions] protocols: NextDnsUpdateCoordinator[AnalyticsProtocols] settings: NextDnsUpdateCoordinator[Settings] status: NextDnsUpdateCoordinator[AnalyticsStatus] @dataclass class NextDnsData: """Runtime data for the NextDNS integration.""" client: NextDns profiles: dict[str, NextDnsCoordinators] CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.SENSOR, Platform.SWITCH] COORDINATORS: list[tuple[str, type[NextDnsUpdateCoordinator]]] = [ (ATTR_CONNECTION, NextDnsConnectionUpdateCoordinator), (ATTR_DNSSEC, NextDnsDnssecUpdateCoordinator), (ATTR_ENCRYPTION, NextDnsEncryptionUpdateCoordinator), (ATTR_IP_VERSIONS, NextDnsIpVersionsUpdateCoordinator), (ATTR_PROTOCOLS, NextDnsProtocolsUpdateCoordinator), (ATTR_SETTINGS, NextDnsSettingsUpdateCoordinator), (ATTR_STATUS, NextDnsStatusUpdateCoordinator), ] async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up NextDNS.""" await async_migrate_integration(hass) return True async def async_migrate_integration(hass: HomeAssistant) -> None: """Migrate integration entry structure.""" # Make sure we get enabled config entries first entries = sorted( hass.config_entries.async_entries(DOMAIN), key=lambda e: e.disabled_by is not None, ) if not any(entry.version == 1 for entry in entries): return api_keys_entries: dict[str, tuple[NextDnsConfigEntry, bool]] = {} device_registry = dr.async_get(hass) entity_registry = er.async_get(hass) for entry in entries: profile_id = entry.data[CONF_PROFILE_ID] profile_name = entry.title subentry = ConfigSubentry( data=MappingProxyType({CONF_PROFILE_ID: profile_id}), subentry_type=SUBENTRY_TYPE_PROFILE, title=profile_name, unique_id=profile_id, ) if entry.data[CONF_API_KEY] not in api_keys_entries: all_disabled = all( e.disabled_by is not None for e in entries if e.data[CONF_API_KEY] == entry.data[CONF_API_KEY] ) api_keys_entries[entry.data[CONF_API_KEY]] = (entry, all_disabled) parent_entry, all_disabled = api_keys_entries[entry.data[CONF_API_KEY]] hass.config_entries.async_add_subentry(parent_entry, subentry) entities = er.async_entries_for_config_entry(entity_registry, entry.entry_id) device = device_registry.async_get_device_by_identifier( (DOMAIN, profile_id), entry.entry_id ) for entity_entry in entities: entity_disabled_by = entity_entry.disabled_by if ( entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY and not all_disabled ): # Device and entity registries don't update the disabled_by flag # when moving a device or entity from one config entry to another, # so we need to do it manually. entity_disabled_by = ( er.RegistryEntryDisabler.DEVICE if device else er.RegistryEntryDisabler.USER ) entity_registry.async_update_entity( entity_entry.entity_id, config_entry_id=parent_entry.entry_id, config_subentry_id=subentry.subentry_id, disabled_by=entity_disabled_by, ) if device is not None: # Device and entity registries don't update the disabled_by flag when # moving a device or entity from one config entry to another, so we # need to do it manually. device_disabled_by = device.disabled_by if ( device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY and not all_disabled ): device_disabled_by = dr.DeviceEntryDisabler.USER device_registry.async_update_device( device.id, disabled_by=device_disabled_by, new_identifiers={(DOMAIN, profile_id)}, new_config_entry_id=parent_entry.entry_id, new_config_subentry_id=subentry.subentry_id, ) if parent_entry.entry_id != entry.entry_id: await hass.config_entries.async_remove(entry.entry_id) else: hass.config_entries.async_update_entry( entry, data={CONF_API_KEY: entry.data[CONF_API_KEY]}, title="NextDNS", version=2, unique_id=None, ) async def async_setup_entry(hass: HomeAssistant, entry: NextDnsConfigEntry) -> bool: """Set up NextDNS as config entry.""" api_key = entry.data[CONF_API_KEY] websession = async_get_clientsession(hass) try: nextdns = await NextDns.create(websession, api_key) except (ApiError, ClientConnectorError, RetryError, TimeoutError) as err: raise ConfigEntryNotReady( translation_domain=DOMAIN, translation_key="cannot_connect", translation_placeholders={ "entry": entry.title, "error": repr(err), }, ) from err except InvalidApiKeyError as err: raise ConfigEntryAuthFailed( translation_domain=DOMAIN, translation_key="auth_error", translation_placeholders={"entry": entry.title}, ) from err profiles: dict[str, NextDnsCoordinators] = {} for subentry in entry.get_subentries_of_type(SUBENTRY_TYPE_PROFILE): subentry_id = subentry.subentry_id profile_id = subentry.data[CONF_PROFILE_ID] tasks = [] coordinators = {} # Independent DataUpdateCoordinator is used for each API endpoint to avoid # unnecessary requests when entities using this endpoint are disabled. for coordinator_name, coordinator_class in COORDINATORS: coordinator = coordinator_class( hass, entry, nextdns, profile_id, subentry_id ) tasks.append(coordinator.async_config_entry_first_refresh()) coordinators[coordinator_name] = coordinator await asyncio.gather(*tasks) profiles[subentry_id] = NextDnsCoordinators(**coordinators) entry.runtime_data = NextDnsData(client=nextdns, profiles=profiles) entry.async_on_unload(entry.add_update_listener(_async_update_listener)) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True async def _async_update_listener( hass: HomeAssistant, entry: NextDnsConfigEntry ) -> None: """Reload the config entry when subentries change.""" await hass.config_entries.async_reload(entry.entry_id) async def async_unload_entry(hass: HomeAssistant, entry: NextDnsConfigEntry) -> bool: """Unload a config entry.""" return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)