config_flow.py

Go to the documentation of this file.

"""Config flow to configure roomba component."""
 
from __future__ import annotations
 
import asyncio
from functools import partial
from typing import Any
 
from roombapy import RoombaFactory, RoombaInfo
from roombapy.discovery import RoombaDiscovery
from roombapy.getpassword import RoombaPassword
import voluptuous as vol
 
from homeassistant.components import dhcp, zeroconf
from homeassistant.config_entries import (
    ConfigEntry,
    ConfigFlow,
    ConfigFlowResult,
    OptionsFlow,
)
from homeassistant.const import CONF_DELAY, CONF_HOST, CONF_NAME, CONF_PASSWORD
from homeassistant.core import HomeAssistant, callback
 
from . import CannotConnect, async_connect_or_timeout, async_disconnect_or_timeout
from .const import (
    CONF_BLID,
    CONF_CONTINUOUS,
    DEFAULT_CONTINUOUS,
    DEFAULT_DELAY,
    DOMAIN,
    ROOMBA_SESSION,
)
 
ROOMBA_DISCOVERY_LOCK = "roomba_discovery_lock"
ALL_ATTEMPTS = 2
HOST_ATTEMPTS = 6
ROOMBA_WAKE_TIME = 6
 
DEFAULT_OPTIONS = {CONF_CONTINUOUS: DEFAULT_CONTINUOUS, CONF_DELAY: DEFAULT_DELAY}
 
MAX_NUM_DEVICES_TO_DISCOVER = 25
 
AUTH_HELP_URL_KEY = "auth_help_url"
AUTH_HELP_URL_VALUE = (
    "https://www.home-assistant.io/integrations/roomba/#retrieving-your-credentials"
)
 
 
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.
    """
    roomba = await hass.async_add_executor_job(
        partial(
            RoombaFactory.create_roomba,
            address=data[CONF_HOST],
            blid=data[CONF_BLID],
            password=data[CONF_PASSWORD],
            continuous=True,
            delay=data[CONF_DELAY],
        )
    )
 
    info = await async_connect_or_timeout(hass, roomba)
    if info:
        await async_disconnect_or_timeout(hass, roomba)
 
    return {
        ROOMBA_SESSION: info[ROOMBA_SESSION],
        CONF_NAME: info[CONF_NAME],
        CONF_HOST: data[CONF_HOST],
    }
 
 
class RoombaConfigFlow(ConfigFlow, domain=DOMAIN):
    """Roomba configuration flow."""
 
    VERSION = 1
 
    name: str | None = None
    blid: str
    host: str | None = None
 
    def __init__(self) -> None:
        """Initialize the roomba flow."""
        self.discovered_robotsdiscovered_robots: dict[str, RoombaInfo] = {}
 
    @staticmethod
    @callback
    def async_get_options_flow(
        config_entry: ConfigEntry,
    ) -> RoombaOptionsFlowHandler:
        """Get the options flow for this handler."""
        return RoombaOptionsFlowHandler()
 
    async def async_step_zeroconf(
        self, discovery_info: zeroconf.ZeroconfServiceInfo
    ) -> ConfigFlowResult:
        """Handle zeroconf discovery."""
        return await self._async_step_discovery_async_step_discovery(
            discovery_info.host, discovery_info.hostname.lower().removesuffix(".local.")
        )
 
    async def async_step_dhcp(
        self, discovery_info: dhcp.DhcpServiceInfo
    ) -> ConfigFlowResult:
        """Handle dhcp discovery."""
        return await self._async_step_discovery_async_step_discovery(
            discovery_info.ip, discovery_info.hostname
        )
 
    async def _async_step_discovery(
        self, ip_address: str, hostname: str
    ) -> ConfigFlowResult:
        """Handle any discovery."""
        self._async_abort_entries_match_async_abort_entries_match({CONF_HOST: ip_address})
 
        if not hostname.startswith(("irobot-", "roomba-")):
            return self.async_abortasync_abortasync_abort(reason="not_irobot_device")
 
        self.hosthost = ip_address
        self.blidblid = _async_blid_from_hostname(hostname)
        await self.async_set_unique_idasync_set_unique_id(self.blidblid)
        self._abort_if_unique_id_configured_abort_if_unique_id_configured(updates={CONF_HOST: ip_address})
 
        # Because the hostname is so long some sources may
        # truncate the hostname since it will be longer than
        # the valid allowed length. If we already have a flow
        # going for a longer hostname we abort so the user
        # does not see two flows if discovery fails.
        for progress in self._async_in_progress_async_in_progress():
            flow_unique_id = progress["context"].get("unique_id")
            if not flow_unique_id:
                continue
            if flow_unique_id.startswith(self.blidblid):
                return self.async_abortasync_abortasync_abort(reason="short_blid")
            if self.blidblid.startswith(flow_unique_id):
                self.hass.config_entries.flow.async_abort(progress["flow_id"])
 
        self.context["title_placeholders"] = {"host": self.hosthost, "name": self.blidblid}
        return await self.async_step_userasync_step_userasync_step_user()
 
    async def _async_start_link(self) -> ConfigFlowResult:
        """Start linking."""
        assert self.hosthost
        device = self.discovered_robotsdiscovered_robots[self.hosthost]
        self.blidblid = device.blid
        self.namename = device.robot_name
        await self.async_set_unique_idasync_set_unique_id(self.blidblid, raise_on_progress=False)
        self._abort_if_unique_id_configured_abort_if_unique_id_configured()
        return await self.async_step_linkasync_step_link()
 
    async def async_step_user(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle a flow start."""
        # Check if user chooses manual entry
        if user_input is not None and not user_input.get(CONF_HOST):
            return await self.async_step_manualasync_step_manual()
 
        if (
            user_input is not None
            and self.discovered_robotsdiscovered_robots is not None
            and user_input[CONF_HOST] in self.discovered_robotsdiscovered_robots
        ):
            self.hosthost = user_input[CONF_HOST]
            return await self._async_start_link_async_start_link()
 
        already_configured = self._async_current_ids_async_current_ids(False)
 
        devices = await _async_discover_roombas(self.hass, self.hosthost)
 
        if devices:
            # Find already configured hosts
            self.discovered_robotsdiscovered_robots = {
                device.ip: device
                for device in devices
                if device.blid not in already_configured
            }
 
        if self.hosthost and self.hosthost in self.discovered_robotsdiscovered_robots:
            # From discovery
            self.context["title_placeholders"] = {
                "host": self.hosthost,
                "name": self.discovered_robotsdiscovered_robots[self.hosthost].robot_name,
            }
            return await self._async_start_link_async_start_link()
 
        if not self.discovered_robotsdiscovered_robots:
            return await self.async_step_manualasync_step_manual()
 
        hosts: dict[str | None, str] = {
            **{
                device.ip: f"{device.robot_name} ({device.ip})"
                for device in devices
                if device.blid not in already_configured
            },
            None: "Manually add a Roomba or Braava",
        }
 
        return self.async_show_formasync_show_formasync_show_form(
            step_id="user",
            data_schema=vol.Schema({vol.Optional("host"): vol.In(hosts)}),
        )
 
    async def async_step_manual(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle manual device setup."""
        if user_input is None:
            return self.async_show_formasync_show_formasync_show_form(
                step_id="manual",
                description_placeholders={AUTH_HELP_URL_KEY: AUTH_HELP_URL_VALUE},
                data_schema=vol.Schema(
                    {vol.Required(CONF_HOST, default=self.hosthost): str}
                ),
            )
 
        self._async_abort_entries_match_async_abort_entries_match({CONF_HOST: user_input["host"]})
 
        self.hosthost = user_input[CONF_HOST]
 
        devices = await _async_discover_roombas(self.hass, self.hosthost)
        if not devices:
            return self.async_abortasync_abortasync_abort(reason="cannot_connect")
        self.blidblid = devices[0].blid
        self.namename = devices[0].robot_name
 
        await self.async_set_unique_idasync_set_unique_id(self.blidblid, raise_on_progress=False)
        self._abort_if_unique_id_configured_abort_if_unique_id_configured()
        return await self.async_step_linkasync_step_link()
 
    async def async_step_link(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Attempt to link with the Roomba.
 
        Given a configured host, will ask the user to press the home and target buttons
        to connect to the device.
        """
        if user_input is None:
            return self.async_show_formasync_show_formasync_show_form(
                step_id="link",
                description_placeholders={CONF_NAME: self.namename or self.blidblid},
            )
        assert self.hosthost
        roomba_pw = RoombaPassword(self.hosthost)
 
        try:
            password = await self.hass.async_add_executor_job(roomba_pw.get_password)
        except OSError:
            return await self.async_step_link_manualasync_step_link_manual()
 
        if not password:
            return await self.async_step_link_manualasync_step_link_manual()
 
        config = {
            CONF_HOST: self.hosthost,
            CONF_BLID: self.blidblid,
            CONF_PASSWORD: password,
            **DEFAULT_OPTIONS,
        }
 
        if not self.namename:
            try:
                info = await validate_input(self.hass, config)
            except CannotConnect:
                return self.async_abortasync_abortasync_abort(reason="cannot_connect")
 
            self.namename = info[CONF_NAME]
        assert self.namename
        return self.async_create_entryasync_create_entryasync_create_entry(title=self.namename, data=config)
 
    async def async_step_link_manual(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Handle manual linking."""
        errors = {}
 
        if user_input is not None:
            config = {
                CONF_HOST: self.hosthost,
                CONF_BLID: self.blidblid,
                CONF_PASSWORD: user_input[CONF_PASSWORD],
                **DEFAULT_OPTIONS,
            }
            try:
                info = await validate_input(self.hass, config)
            except CannotConnect:
                errors = {"base": "cannot_connect"}
            else:
                return self.async_create_entryasync_create_entryasync_create_entry(title=info[CONF_NAME], data=config)
 
        return self.async_show_formasync_show_formasync_show_form(
            step_id="link_manual",
            description_placeholders={AUTH_HELP_URL_KEY: AUTH_HELP_URL_VALUE},
            data_schema=vol.Schema({vol.Required(CONF_PASSWORD): str}),
            errors=errors,
        )
 
 
class RoombaOptionsFlowHandler(OptionsFlow):
    """Handle options."""
 
    async def async_step_init(
        self, user_input: dict[str, Any] | None = None
    ) -> ConfigFlowResult:
        """Manage the options."""
        if user_input is not None:
            return self.async_create_entryasync_create_entry(title="", data=user_input)
 
        options = self.config_entryconfig_entryconfig_entry.options
        return self.async_show_formasync_show_form(
            step_id="init",
            data_schema=vol.Schema(
                {
                    vol.Optional(
                        CONF_CONTINUOUS,
                        default=options.get(CONF_CONTINUOUS, DEFAULT_CONTINUOUS),
                    ): bool,
                    vol.Optional(
                        CONF_DELAY,
                        default=options.get(CONF_DELAY, DEFAULT_DELAY),
                    ): int,
                }
            ),
        )
 
 
@callback
def _async_get_roomba_discovery() -> RoombaDiscovery:
    """Create a discovery object."""
    discovery = RoombaDiscovery()
    discovery.amount_of_broadcasted_messages = MAX_NUM_DEVICES_TO_DISCOVER
    return discovery
 
 
@callback
def _async_blid_from_hostname(hostname: str) -> str:
    """Extract the blid from the hostname."""
    return hostname.split("-")[1].split(".")[0].upper()
 
 
async def _async_discover_roombas(
    hass: HomeAssistant, host: str | None = None
) -> list[RoombaInfo]:
    discovered_hosts: set[str] = set()
    devices: list[RoombaInfo] = []
    discover_lock = hass.data.setdefault(ROOMBA_DISCOVERY_LOCK, asyncio.Lock())
    discover_attempts = HOST_ATTEMPTS if host else ALL_ATTEMPTS
 
    for attempt in range(discover_attempts + 1):
        async with discover_lock:
            discovery = _async_get_roomba_discovery()
            discovered: set[RoombaInfo] = set()
            try:
                if host:
                    device = await hass.async_add_executor_job(discovery.get, host)
                    if device:
                        discovered.add(device)
                else:
                    discovered = await hass.async_add_executor_job(discovery.get_all)
            except OSError:
                # Socket temporarily unavailable
                await asyncio.sleep(ROOMBA_WAKE_TIME * attempt)
                continue
            else:
                for device in discovered:
                    if device.ip in discovered_hosts:
                        continue
                    discovered_hosts.add(device.ip)
                    devices.append(device)
            finally:
                discovery.server_socket.close()
 
        if host and host in discovered_hosts:
            return devices
 
        await asyncio.sleep(ROOMBA_WAKE_TIME)
 
    return devices