Add config flow to imap (home-assistant#74623)

* Add config flow to imap

fix coverage

fix config_flows.py

* move coordinator to seperate file, remove name key

* update intrgations.json

* update requirements_all.txt

* fix importing issue_registry

* Address comments

* Improve handling exceptions on intial connection

* exit loop tasks properly

* fix timeout

* revert async_timeout

* Improve entity update handling

* ensure we wait for idle to finish

* fix typing

* Update deprecation period

Co-authored-by: Martin Hjelmare <[email protected]>

.coveragerc

@@ -561,6 +561,8 @@ omit =
     homeassistant/components/ifttt/const.py
     homeassistant/components/iglo/light.py
     homeassistant/components/ihc/*
+    homeassistant/components/imap/__init__.py
+    homeassistant/components/imap/coordinator.py
     homeassistant/components/imap/sensor.py
     homeassistant/components/imap_email_content/sensor.py
     homeassistant/components/incomfort/*

CODEOWNERS

@@ -537,6 +537,8 @@ build.json @home-assistant/supervisor
 /tests/components/image_processing/ @home-assistant/core
 /homeassistant/components/image_upload/ @home-assistant/core
 /tests/components/image_upload/ @home-assistant/core
+/homeassistant/components/imap/ @engrbm87
+/tests/components/imap/ @engrbm87
 /homeassistant/components/incomfort/ @zxdavb
 /homeassistant/components/influxdb/ @mdegat01
 /tests/components/influxdb/ @mdegat01

homeassistant/components/imap/__init__.py

@@ -1 +1,54 @@
-"""The imap component."""
+"""The imap integration."""
+from __future__ import annotations
+
+import asyncio
+
+from aioimaplib import IMAP4_SSL, AioImapException
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import (
+    ConfigEntryAuthFailed,
+    ConfigEntryError,
+    ConfigEntryNotReady,
+)
+
+from .const import DOMAIN
+from .coordinator import ImapDataUpdateCoordinator, connect_to_server
+from .errors import InvalidAuth, InvalidFolder
+
+PLATFORMS: list[Platform] = [Platform.SENSOR]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Set up imap from a config entry."""
+    try:
+        imap_client: IMAP4_SSL = await connect_to_server(dict(entry.data))
+    except InvalidAuth as err:
+        raise ConfigEntryAuthFailed from err
+    except InvalidFolder as err:
+        raise ConfigEntryError("Selected mailbox folder is invalid.") from err
+    except (asyncio.TimeoutError, AioImapException) as err:
+        raise ConfigEntryNotReady from err
+
+    coordinator = ImapDataUpdateCoordinator(hass, imap_client)
+    await coordinator.async_config_entry_first_refresh()
+
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, coordinator.shutdown)
+    )
+
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Unload a config entry."""
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+        coordinator: ImapDataUpdateCoordinator = hass.data[DOMAIN].pop(entry.entry_id)
+        await coordinator.shutdown()
+    return unload_ok

homeassistant/components/imap/config_flow.py

@@ -0,0 +1,136 @@
+"""Config flow for imap integration."""
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Mapping
+from typing import Any
+
+from aioimaplib import AioImapException
+import voluptuous as vol
+
+from homeassistant import config_entries
+from homeassistant.const import CONF_NAME, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
+from homeassistant.data_entry_flow import FlowResult
+from homeassistant.helpers import config_validation as cv
+
+from .const import (
+    CONF_CHARSET,
+    CONF_FOLDER,
+    CONF_SEARCH,
+    CONF_SERVER,
+    DEFAULT_PORT,
+    DOMAIN,
+)
+from .coordinator import connect_to_server
+from .errors import InvalidAuth, InvalidFolder
+
+STEP_USER_DATA_SCHEMA = vol.Schema(
+    {
+        vol.Required(CONF_USERNAME): str,
+        vol.Required(CONF_PASSWORD): str,
+        vol.Required(CONF_SERVER): str,
+        vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
+        vol.Optional(CONF_CHARSET, default="utf-8"): str,
+        vol.Optional(CONF_FOLDER, default="INBOX"): str,
+        vol.Optional(CONF_SEARCH, default="UnSeen UnDeleted"): str,
+    }
+)
+
+
+async def validate_input(user_input: dict[str, Any]) -> dict[str, str]:
+    """Validate user input."""
+    errors = {}
+
+    try:
+        imap_client = await connect_to_server(user_input)
+        result, lines = await imap_client.search(
+            user_input[CONF_SEARCH],
+            charset=user_input[CONF_CHARSET],
+        )
+
+    except InvalidAuth:
+        errors[CONF_USERNAME] = errors[CONF_PASSWORD] = "invalid_auth"
+    except InvalidFolder:
+        errors[CONF_FOLDER] = "invalid_folder"
+    except (asyncio.TimeoutError, AioImapException, ConnectionRefusedError):
+        errors["base"] = "cannot_connect"
+    else:
+        if result != "OK":
+            if "The specified charset is not supported" in lines[0].decode("utf-8"):
+                errors[CONF_CHARSET] = "invalid_charset"
+            else:
+                errors[CONF_SEARCH] = "invalid_search"
+    return errors
+
+
+class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for imap."""
+
+    VERSION = 1
+    _reauth_entry: config_entries.ConfigEntry | None
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle the initial step."""
+        if user_input is None:
+            return self.async_show_form(
+                step_id="user", data_schema=STEP_USER_DATA_SCHEMA
+            )
+
+        self._async_abort_entries_match(
+            {
+                CONF_USERNAME: user_input[CONF_USERNAME],
+                CONF_FOLDER: user_input[CONF_FOLDER],
+                CONF_SEARCH: user_input[CONF_SEARCH],
+            }
+        )
+
+        if not (errors := await validate_input(user_input)):
+            # To be removed when YAML import is removed
+            title = user_input.get(CONF_NAME, user_input[CONF_USERNAME])
+
+            return self.async_create_entry(title=title, data=user_input)
+
+        return self.async_show_form(
+            step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
+        )
+
+    async def async_step_import(self, import_config: dict[str, Any]) -> FlowResult:
+        """Import a config entry from configuration.yaml."""
+        return await self.async_step_user(import_config)
+
+    async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
+        """Perform reauth upon an API authentication error."""
+        self._reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+        return await self.async_step_reauth_confirm()
+
+    async def async_step_reauth_confirm(
+        self, user_input: dict[str, str] | None = None
+    ) -> FlowResult:
+        """Confirm reauth dialog."""
+        errors = {}
+        assert self._reauth_entry
+        if user_input is not None:
+            user_input = {**self._reauth_entry.data, **user_input}
+            if not (errors := await validate_input(user_input)):
+                self.hass.config_entries.async_update_entry(
+                    self._reauth_entry, data=user_input
+                )
+                await self.hass.config_entries.async_reload(self._reauth_entry.entry_id)
+                return self.async_abort(reason="reauth_successful")
+
+        return self.async_show_form(
+            description_placeholders={
+                CONF_USERNAME: self._reauth_entry.data[CONF_USERNAME]
+            },
+            step_id="reauth_confirm",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_PASSWORD): str,
+                }
+            ),
+            errors=errors,
+        )

homeassistant/components/imap/const.py

@@ -0,0 +1,12 @@
+"""Constants for the imap integration."""
+
+from typing import Final
+
+DOMAIN: Final = "imap"
+
+CONF_SERVER: Final = "server"
+CONF_FOLDER: Final = "folder"
+CONF_SEARCH: Final = "search"
+CONF_CHARSET: Final = "charset"
+
+DEFAULT_PORT: Final = 993

homeassistant/components/imap/coordinator.py

@@ -0,0 +1,104 @@
+"""Coordinator for imag integration."""
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Mapping
+from datetime import timedelta
+import logging
+from typing import Any
+
+from aioimaplib import AUTH, IMAP4_SSL, SELECTED, AioImapException
+import async_timeout
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_PASSWORD, CONF_PORT, CONF_USERNAME
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import CONF_CHARSET, CONF_FOLDER, CONF_SEARCH, CONF_SERVER, DOMAIN
+from .errors import InvalidAuth, InvalidFolder
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def connect_to_server(data: Mapping[str, Any]) -> IMAP4_SSL:
+    """Connect to imap server and return client."""
+    client = IMAP4_SSL(data[CONF_SERVER], data[CONF_PORT])
+    await client.wait_hello_from_server()
+    await client.login(data[CONF_USERNAME], data[CONF_PASSWORD])
+    if client.protocol.state != AUTH:
+        raise InvalidAuth
+    await client.select(data[CONF_FOLDER])
+    if client.protocol.state != SELECTED:
+        raise InvalidFolder
+    return client
+
+
+class ImapDataUpdateCoordinator(DataUpdateCoordinator[int]):
+    """Class for imap client."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, imap_client: IMAP4_SSL) -> None:
+        """Initiate imap client."""
+        self.hass = hass
+        self.imap_client = imap_client
+        self.support_push = imap_client.has_capability("IDLE")
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=timedelta(seconds=10) if not self.support_push else None,
+        )
+
+    async def _async_update_data(self) -> int:
+        """Update the number of unread emails."""
+        try:
+            if self.imap_client is None:
+                self.imap_client = await connect_to_server(self.config_entry.data)
+        except (AioImapException, asyncio.TimeoutError) as err:
+            raise UpdateFailed(err) from err
+
+        return await self.refresh_email_count()
+
+    async def refresh_email_count(self) -> int:
+        """Check the number of found emails."""
+        try:
+            await self.imap_client.noop()
+            result, lines = await self.imap_client.search(
+                self.config_entry.data[CONF_SEARCH],
+                charset=self.config_entry.data[CONF_CHARSET],
+            )
+        except (AioImapException, asyncio.TimeoutError) as err:
+            raise UpdateFailed(err) from err
+
+        if result != "OK":
+            raise UpdateFailed(
+                f"Invalid response for search '{self.config_entry.data[CONF_SEARCH]}': {result} / {lines[0]}"
+            )
+        if self.support_push:
+            self.hass.async_create_task(self.async_wait_server_push())
+        return len(lines[0].split())
+
+    async def async_wait_server_push(self) -> None:
+        """Wait for data push from server."""
+        try:
+            idle: asyncio.Future = await self.imap_client.idle_start()
+            await self.imap_client.wait_server_push()
+            self.imap_client.idle_done()
+            async with async_timeout.timeout(10):
+                await idle
+
+        except (AioImapException, asyncio.TimeoutError):
+            _LOGGER.warning(
+                "Lost %s (will attempt to reconnect)",
+                self.config_entry.data[CONF_SERVER],
+            )
+            self.imap_client = None
+        await self.async_request_refresh()
+
+    async def shutdown(self, *_) -> None:
+        """Close resources."""
+        if self.imap_client:
+            await self.imap_client.stop_wait_server_push()
+            await self.imap_client.logout()

homeassistant/components/imap/errors.py

@@ -0,0 +1,11 @@
+"""Exceptions raised by IMAP integration."""
+
+from homeassistant.exceptions import HomeAssistantError
+
+
+class InvalidAuth(HomeAssistantError):
+    """Raise exception for invalid credentials."""
+
+
+class InvalidFolder(HomeAssistantError):
+    """Raise exception for invalid folder."""

homeassistant/components/imap/manifest.json

@@ -1,9 +1,11 @@
 {
   "domain": "imap",
   "name": "IMAP",
+  "config_flow": true,
+  "dependencies": ["repairs"],
   "documentation": "https://www.home-assistant.io/integrations/imap",
   "requirements": ["aioimaplib==1.0.1"],
-  "codeowners": [],
+  "codeowners": ["@engrbm87"],
   "iot_class": "cloud_push",
   "loggers": ["aioimaplib"]
 }