config_flow.py

Go to the documentation of this file.

"""Config flow for Cloudflare integration."""
 
from __future__ import annotations
 
from collections.abc import Mapping
import logging
from typing import Any
 
import pycfdns
import voluptuous as vol
 
from homeassistant.components import persistent_notification
from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
from homeassistant.const import CONF_API_TOKEN, CONF_ZONE
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import config_validation as cv
from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
from .const import CONF_RECORDS, DOMAIN
from .helpers import get_zone_id
 
_LOGGER = logging.getLogger(__name__)
 
DATA_SCHEMA = vol.Schema(
    {
        vol.Required(CONF_API_TOKEN): str,
    }
)
 
 
def _zone_schema(zones: list[pycfdns.ZoneModel] | None = None) -> vol.Schema:
    """Zone selection schema."""
    zones_list = []
 
    if zones is not None:
        zones_list = [zones["name"] for zones in zones]
 
    return vol.Schema({vol.Required(CONF_ZONE): vol.In(zones_list)})
 
 
def _records_schema(records: list[pycfdns.RecordModel] | None = None) -> vol.Schema:
    """Zone records selection schema."""
    records_dict = {}
 
    if records:
        records_dict = {name["name"]: name["name"] for name in records}
 
    return vol.Schema({vol.Required(CONF_RECORDS): cv.multi_select(records_dict)})
 
 
async def _validate_input(
    hass: HomeAssistant,
    data: dict[str, Any],
) -> dict[str, Any]:
    """Validate the user input allows us to connect.
 
    Data has the keys from DATA_SCHEMA with values provided by the user.
    """
    zone = data.get(CONF_ZONE)
    records: list[pycfdns.RecordModel] = []
 
    client = pycfdns.Client(
        api_token=data[CONF_API_TOKEN],
        client_session=async_get_clientsession(hass),
    )
 
    zones = await client.list_zones()
    if zone and (zone_id := get_zone_id(zone, zones)) is not None:
        records = await client.list_dns_records(zone_id=zone_id, type="A")
 
    return {"zones": zones, "records": records}
 
 
class CloudflareConfigFlow(ConfigFlow, domain=DOMAIN):
    """Handle a config flow for Cloudflare."""
 
    VERSION = 1
 
    def __init__(self) -> None:
        """Initialize the Cloudflare config flow."""
        self.cloudflare_config: dict[str, Any] = {}
        self.zoneszones: list[pycfdns.ZoneModel] | None = None
        self.recordsrecords: list[pycfdns.RecordModel] | None = None
 
    async def async_step_reauth(
        self, entry_data: Mapping[str, Any]
    ) -> ConfigFlowResult:
        """Handle initiation of re-authentication with Cloudflare."""
        return await self.async_step_reauth_confirmasync_step_reauth_confirm()
 
    async def async_step_reauth_confirm(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle re-authentication with Cloudflare."""
        errors: dict[str, str] = {}
 
        if user_input is not None:
            _, errors = await self._async_validate_or_error_async_validate_or_error(user_input)
 
            if not errors:
                reauth_entry = self._get_reauth_entry_get_reauth_entry()
                return self.async_update_reload_and_abortasync_update_reload_and_abort(
                    reauth_entry,
                    data={
                        **reauth_entry.data,
                        CONF_API_TOKEN: user_input[CONF_API_TOKEN],
                    },
                )
 
        return self.async_show_formasync_show_formasync_show_form(
            step_id="reauth_confirm",
            data_schema=DATA_SCHEMA,
            errors=errors,
        )
 
    async def async_step_user(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle a flow initiated by the user."""
        persistent_notification.async_dismiss(self.hass, "cloudflare_setup")
 
        errors: dict[str, str] = {}
 
        if user_input is not None:
            info, errors = await self._async_validate_or_error_async_validate_or_error(user_input)
 
            if not errors:
                self.cloudflare_config.update(user_input)
                self.zoneszones = info["zones"]
                return await self.async_step_zoneasync_step_zone()
 
        return self.async_show_formasync_show_formasync_show_form(
            step_id="user", data_schema=DATA_SCHEMA, errors=errors
        )
 
    async def async_step_zone(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle the picking the zone."""
        errors: dict[str, str] = {}
 
        if user_input is not None:
            self.cloudflare_config.update(user_input)
            info, errors = await self._async_validate_or_error_async_validate_or_error(self.cloudflare_config)
 
            if not errors:
                await self.async_set_unique_idasync_set_unique_id(user_input[CONF_ZONE])
                self.recordsrecords = info["records"]
 
                return await self.async_step_recordsasync_step_records()
 
        return self.async_show_formasync_show_formasync_show_form(
            step_id="zone",
            data_schema=_zone_schema(self.zoneszones),
            errors=errors,
        )
 
    async def async_step_records(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle the picking the zone records."""
 
        if user_input is not None:
            self.cloudflare_config.update(user_input)
            title = self.cloudflare_config[CONF_ZONE]
            return self.async_create_entryasync_create_entryasync_create_entry(title=title, data=self.cloudflare_config)
 
        return self.async_show_formasync_show_formasync_show_form(
            step_id="records",
            data_schema=_records_schema(self.recordsrecords),
        )
 
    async def _async_validate_or_error(
        self, config: dict[str, Any]
    ) -> tuple[dict[str, list[Any]], dict[str, str]]:
        errors: dict[str, str] = {}
        info = {}
 
        try:
            info = await _validate_input(self.hass, config)
        except pycfdns.ComunicationException:
            errors["base"] = "cannot_connect"
        except pycfdns.AuthenticationException:
            errors["base"] = "invalid_auth"
        except Exception:
            _LOGGER.exception("Unexpected exception")
            errors["base"] = "unknown"
 
        return info, errors
 
 
class CannotConnect(HomeAssistantError):
    """Error to indicate we cannot connect."""
 
 
class InvalidAuth(HomeAssistantError):
    """Error to indicate there is invalid auth."""