diff --git a/homeassistant/components/imap/__init__.py b/homeassistant/components/imap/__init__.py index 0ca5384f7e92..c2369935d5fc 100644 --- a/homeassistant/components/imap/__init__.py +++ b/homeassistant/components/imap/__init__.py @@ -1,27 +1,16 @@ """The imap integration.""" -import asyncio -from email.message import Message import logging -from typing import Any -from aioimaplib import IMAP4_SSL, AioImapException, Response -import probatio +from aioimaplib import IMAP4_SSL, AioImapException -from homeassistant.config_entries import ConfigEntry, ConfigEntryState +from homeassistant.config_entries import ConfigEntry from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform -from homeassistant.core import ( - HomeAssistant, - ServiceCall, - ServiceResponse, - SupportsResponse, - callback, -) +from homeassistant.core import HomeAssistant from homeassistant.exceptions import ( ConfigEntryAuthFailed, ConfigEntryError, ConfigEntryNotReady, - ServiceValidationError, ) from homeassistant.helpers import config_validation as cv from homeassistant.helpers.typing import ConfigType @@ -29,287 +18,28 @@ from homeassistant.helpers.typing import ConfigType from .const import CONF_ENABLE_PUSH, DOMAIN from .coordinator import ( ImapDataUpdateCoordinator, - ImapMessage, ImapPollingDataUpdateCoordinator, ImapPushDataUpdateCoordinator, connect_to_server, - get_parts, ) from .errors import InvalidAuth, InvalidFolder +from .services import async_setup_services PLATFORMS: list[Platform] = [Platform.SENSOR] -CONF_ENTRY = "entry" -CONF_SEEN = "seen" -CONF_PART = "part" -CONF_UID = "uid" -CONF_TARGET_FOLDER = "target_folder" _LOGGER = logging.getLogger(__name__) CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) -_SERVICE_UID_SCHEMA = probatio.Schema( - { - probatio.Required(CONF_ENTRY): cv.string, - probatio.Required(CONF_UID): cv.string, - } -) - -SERVICE_SEEN_SCHEMA = _SERVICE_UID_SCHEMA -SERVICE_MOVE_SCHEMA = _SERVICE_UID_SCHEMA.extend( - { - probatio.Optional(CONF_SEEN): cv.boolean, - probatio.Required(CONF_TARGET_FOLDER): cv.string, - } -) -SERVICE_DELETE_SCHEMA = _SERVICE_UID_SCHEMA -SERVICE_FETCH_TEXT_SCHEMA = _SERVICE_UID_SCHEMA -SERVICE_FETCH_PART_SCHEMA = _SERVICE_UID_SCHEMA.extend( - { - probatio.Required(CONF_PART): cv.string, - } -) type ImapConfigEntry = ConfigEntry[ImapDataUpdateCoordinator] -async def async_get_imap_client(hass: HomeAssistant, entry_id: str) -> IMAP4_SSL: - """Get IMAP client and connect.""" - if (entry := hass.config_entries.async_get_entry(entry_id)) is None or ( - entry.state is not ConfigEntryState.LOADED - ): - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="invalid_entry", - ) - try: - client = await connect_to_server(entry.data) - except InvalidAuth as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, translation_key="invalid_auth" - ) from exc - except InvalidFolder as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, translation_key="invalid_folder" - ) from exc - except (TimeoutError, AioImapException) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="imap_server_fail", - translation_placeholders={"error": str(exc)}, - ) from exc - return client - - -@callback -def raise_on_error(response: Response, translation_key: str) -> None: - """Get error message from response.""" - if response.result != "OK": - error: str = response.lines[0].decode("utf-8") - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key=translation_key, - translation_placeholders={"error": error}, - ) - - async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up imap services.""" - - async def async_seen(call: ServiceCall) -> None: - """Process mark as seen service call.""" - entry_id: str = call.data[CONF_ENTRY] - uid: str = call.data[CONF_UID] - _LOGGER.debug( - "Mark message %s as seen. Entry: %s", - uid, - entry_id, - ) - client = await async_get_imap_client(hass, entry_id) - try: - response = await client.store(uid, "+FLAGS (\\Seen)") - except (TimeoutError, AioImapException) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="imap_server_fail", - translation_placeholders={"error": str(exc)}, - ) from exc - raise_on_error(response, "seen_failed") - await client.close() - - hass.services.async_register(DOMAIN, "seen", async_seen, SERVICE_SEEN_SCHEMA) - - async def async_move(call: ServiceCall) -> None: - """Process move email service call.""" - entry_id: str = call.data[CONF_ENTRY] - uid: str = call.data[CONF_UID] - seen = bool(call.data.get(CONF_SEEN)) - target_folder: str = call.data[CONF_TARGET_FOLDER] - _LOGGER.debug( - "Move message %s to folder %s. Mark as seen: %s. Entry: %s", - uid, - target_folder, - seen, - entry_id, - ) - client = await async_get_imap_client(hass, entry_id) - try: - if seen: - response = await client.store(uid, "+FLAGS (\\Seen)") - raise_on_error(response, "seen_failed") - response = await client.copy(uid, target_folder) - raise_on_error(response, "copy_failed") - response = await client.store(uid, "+FLAGS (\\Deleted)") - raise_on_error(response, "delete_failed") - response = await asyncio.wait_for( - client.protocol.expunge(uid, by_uid=True), client.timeout - ) - raise_on_error(response, "expunge_failed") - except (TimeoutError, AioImapException) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="imap_server_fail", - translation_placeholders={"error": str(exc)}, - ) from exc - await client.close() - - hass.services.async_register(DOMAIN, "move", async_move, SERVICE_MOVE_SCHEMA) - - async def async_delete(call: ServiceCall) -> None: - """Process deleting email service call.""" - entry_id: str = call.data[CONF_ENTRY] - uid: str = call.data[CONF_UID] - _LOGGER.debug( - "Delete message %s. Entry: %s", - uid, - entry_id, - ) - client = await async_get_imap_client(hass, entry_id) - try: - response = await client.store(uid, "+FLAGS (\\Deleted)") - raise_on_error(response, "delete_failed") - response = await asyncio.wait_for( - client.protocol.expunge(uid, by_uid=True), client.timeout - ) - raise_on_error(response, "expunge_failed") - except (TimeoutError, AioImapException) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="imap_server_fail", - translation_placeholders={"error": str(exc)}, - ) from exc - await client.close() - - hass.services.async_register(DOMAIN, "delete", async_delete, SERVICE_DELETE_SCHEMA) - - async def async_fetch(call: ServiceCall) -> ServiceResponse: - """Process fetch email service and return content.""" - entry_id: str = call.data[CONF_ENTRY] - uid: str = call.data[CONF_UID] - _LOGGER.debug( - "Fetch text for message %s. Entry: %s", - uid, - entry_id, - ) - client = await async_get_imap_client(hass, entry_id) - try: - response = await client.fetch(uid, "BODY.PEEK[]") - except (TimeoutError, AioImapException) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="imap_server_fail", - translation_placeholders={"error": str(exc)}, - ) from exc - raise_on_error(response, "fetch_failed") - # Index 1 of of the response lines contains the bytearray with the message data - message = ImapMessage(response.lines[1]) - await client.close() - return { - "text": message.text, - "sender": message.sender, - "subject": message.subject, - "parts": get_parts(message.email_message), - "date": message.date.isoformat() if message.date else None, - "uid": uid, - } - - hass.services.async_register( - DOMAIN, - "fetch", - async_fetch, - SERVICE_FETCH_TEXT_SCHEMA, - supports_response=SupportsResponse.ONLY, - ) - - async def async_fetch_part(call: ServiceCall) -> ServiceResponse: - """Process fetch email part service and return content.""" - - @callback - def get_message_part(message: Message, part_key: str) -> Message: - part: Message | Any = message - for index in part_key.split(","): - sub_parts = part.get_payload() - try: - assert isinstance(sub_parts, list) - part = sub_parts[int(index)] - except (AssertionError, ValueError, IndexError) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="invalid_part_index", - ) from exc - - return part - - entry_id: str = call.data[CONF_ENTRY] - uid: str = call.data[CONF_UID] - part_key: str = call.data[CONF_PART] - _LOGGER.debug( - "Fetch part %s for message %s. Entry: %s", - part_key, - uid, - entry_id, - ) - client = await async_get_imap_client(hass, entry_id) - try: - response = await client.fetch(uid, "BODY.PEEK[]") - except (TimeoutError, AioImapException) as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="imap_server_fail", - translation_placeholders={"error": str(exc)}, - ) from exc - raise_on_error(response, "fetch_failed") - # Index 1 of of the response lines contains the bytearray with the message data - message = ImapMessage(response.lines[1]) - await client.close() - part_data = get_message_part(message.email_message, part_key) - part_data_content = part_data.get_payload(decode=False) - try: - assert isinstance(part_data_content, str) - except AssertionError as exc: - raise ServiceValidationError( - translation_domain=DOMAIN, - translation_key="invalid_part_index", - ) from exc - return { - "part_data": part_data_content, - "content_type": part_data.get_content_type(), - "content_transfer_encoding": part_data.get("Content-Transfer-Encoding"), - "filename": part_data.get_filename(), - "part": part_key, - "uid": uid, - } - - hass.services.async_register( - DOMAIN, - "fetch_part", - async_fetch_part, - SERVICE_FETCH_PART_SCHEMA, - supports_response=SupportsResponse.ONLY, - ) - + async_setup_services(hass) return True diff --git a/homeassistant/components/imap/services.py b/homeassistant/components/imap/services.py new file mode 100644 index 000000000000..91fc674cf966 --- /dev/null +++ b/homeassistant/components/imap/services.py @@ -0,0 +1,291 @@ +"""Support for the imap services.""" + +import asyncio +from email.message import Message +import logging +from typing import Any + +from aioimaplib import IMAP4_SSL, AioImapException, Response +import probatio + +from homeassistant.config_entries import ConfigEntryState +from homeassistant.core import ( + HomeAssistant, + ServiceCall, + ServiceResponse, + SupportsResponse, + callback, +) +from homeassistant.exceptions import ServiceValidationError +from homeassistant.helpers import config_validation as cv + +from .const import DOMAIN +from .coordinator import ImapMessage, connect_to_server, get_parts +from .errors import InvalidAuth, InvalidFolder + +_LOGGER = logging.getLogger(__name__) + +CONF_ENTRY = "entry" +CONF_SEEN = "seen" +CONF_PART = "part" +CONF_UID = "uid" +CONF_TARGET_FOLDER = "target_folder" + +_SERVICE_UID_SCHEMA = probatio.Schema( + { + probatio.Required(CONF_ENTRY): cv.string, + probatio.Required(CONF_UID): cv.string, + } +) + +SERVICE_SEEN_SCHEMA = _SERVICE_UID_SCHEMA +SERVICE_MOVE_SCHEMA = _SERVICE_UID_SCHEMA.extend( + { + probatio.Optional(CONF_SEEN): cv.boolean, + probatio.Required(CONF_TARGET_FOLDER): cv.string, + } +) +SERVICE_DELETE_SCHEMA = _SERVICE_UID_SCHEMA +SERVICE_FETCH_TEXT_SCHEMA = _SERVICE_UID_SCHEMA +SERVICE_FETCH_PART_SCHEMA = _SERVICE_UID_SCHEMA.extend( + { + probatio.Required(CONF_PART): cv.string, + } +) + + +async def async_get_imap_client(hass: HomeAssistant, entry_id: str) -> IMAP4_SSL: + """Get IMAP client and connect.""" + if (entry := hass.config_entries.async_get_entry(entry_id)) is None or ( + entry.state is not ConfigEntryState.LOADED + ): + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="invalid_entry", + ) + try: + client = await connect_to_server(entry.data) + except InvalidAuth as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, translation_key="invalid_auth" + ) from exc + except InvalidFolder as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, translation_key="invalid_folder" + ) from exc + except (TimeoutError, AioImapException) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="imap_server_fail", + translation_placeholders={"error": str(exc)}, + ) from exc + return client + + +@callback +def raise_on_error(response: Response, translation_key: str) -> None: + """Get error message from response.""" + if response.result != "OK": + error: str = response.lines[0].decode("utf-8") + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key=translation_key, + translation_placeholders={"error": error}, + ) + + +@callback +def _get_message_part(message: Message, part_key: str) -> Message: + part: Message | Any = message + for index in part_key.split(","): + sub_parts = part.get_payload() + try: + assert isinstance(sub_parts, list) + part = sub_parts[int(index)] + except (AssertionError, ValueError, IndexError) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="invalid_part_index", + ) from exc + + return part + + +async def _async_seen(call: ServiceCall) -> None: + """Process mark as seen service call.""" + entry_id: str = call.data[CONF_ENTRY] + uid: str = call.data[CONF_UID] + _LOGGER.debug( + "Mark message %s as seen. Entry: %s", + uid, + entry_id, + ) + client = await async_get_imap_client(call.hass, entry_id) + try: + response = await client.store(uid, "+FLAGS (\\Seen)") + except (TimeoutError, AioImapException) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="imap_server_fail", + translation_placeholders={"error": str(exc)}, + ) from exc + raise_on_error(response, "seen_failed") + await client.close() + + +async def _async_move(call: ServiceCall) -> None: + """Process move email service call.""" + entry_id: str = call.data[CONF_ENTRY] + uid: str = call.data[CONF_UID] + seen = bool(call.data.get(CONF_SEEN)) + target_folder: str = call.data[CONF_TARGET_FOLDER] + _LOGGER.debug( + "Move message %s to folder %s. Mark as seen: %s. Entry: %s", + uid, + target_folder, + seen, + entry_id, + ) + client = await async_get_imap_client(call.hass, entry_id) + try: + if seen: + response = await client.store(uid, "+FLAGS (\\Seen)") + raise_on_error(response, "seen_failed") + response = await client.copy(uid, target_folder) + raise_on_error(response, "copy_failed") + response = await client.store(uid, "+FLAGS (\\Deleted)") + raise_on_error(response, "delete_failed") + response = await asyncio.wait_for( + client.protocol.expunge(uid, by_uid=True), client.timeout + ) + raise_on_error(response, "expunge_failed") + except (TimeoutError, AioImapException) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="imap_server_fail", + translation_placeholders={"error": str(exc)}, + ) from exc + await client.close() + + +async def _async_delete(call: ServiceCall) -> None: + """Process deleting email service call.""" + entry_id: str = call.data[CONF_ENTRY] + uid: str = call.data[CONF_UID] + _LOGGER.debug( + "Delete message %s. Entry: %s", + uid, + entry_id, + ) + client = await async_get_imap_client(call.hass, entry_id) + try: + response = await client.store(uid, "+FLAGS (\\Deleted)") + raise_on_error(response, "delete_failed") + response = await asyncio.wait_for( + client.protocol.expunge(uid, by_uid=True), client.timeout + ) + raise_on_error(response, "expunge_failed") + except (TimeoutError, AioImapException) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="imap_server_fail", + translation_placeholders={"error": str(exc)}, + ) from exc + await client.close() + + +async def _async_fetch(call: ServiceCall) -> ServiceResponse: + """Process fetch email service and return content.""" + entry_id: str = call.data[CONF_ENTRY] + uid: str = call.data[CONF_UID] + _LOGGER.debug( + "Fetch text for message %s. Entry: %s", + uid, + entry_id, + ) + client = await async_get_imap_client(call.hass, entry_id) + try: + response = await client.fetch(uid, "BODY.PEEK[]") + except (TimeoutError, AioImapException) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="imap_server_fail", + translation_placeholders={"error": str(exc)}, + ) from exc + raise_on_error(response, "fetch_failed") + # Index 1 of of the response lines contains the bytearray with the message data + message = ImapMessage(response.lines[1]) + await client.close() + return { + "text": message.text, + "sender": message.sender, + "subject": message.subject, + "parts": get_parts(message.email_message), + "date": message.date.isoformat() if message.date else None, + "uid": uid, + } + + +async def _async_fetch_part(call: ServiceCall) -> ServiceResponse: + """Process fetch email part service and return content.""" + entry_id: str = call.data[CONF_ENTRY] + uid: str = call.data[CONF_UID] + part_key: str = call.data[CONF_PART] + _LOGGER.debug( + "Fetch part %s for message %s. Entry: %s", + part_key, + uid, + entry_id, + ) + client = await async_get_imap_client(call.hass, entry_id) + try: + response = await client.fetch(uid, "BODY.PEEK[]") + except (TimeoutError, AioImapException) as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="imap_server_fail", + translation_placeholders={"error": str(exc)}, + ) from exc + raise_on_error(response, "fetch_failed") + # Index 1 of of the response lines contains the bytearray with the message data + message = ImapMessage(response.lines[1]) + await client.close() + part_data = _get_message_part(message.email_message, part_key) + part_data_content = part_data.get_payload(decode=False) + try: + assert isinstance(part_data_content, str) + except AssertionError as exc: + raise ServiceValidationError( + translation_domain=DOMAIN, + translation_key="invalid_part_index", + ) from exc + return { + "part_data": part_data_content, + "content_type": part_data.get_content_type(), + "content_transfer_encoding": part_data.get("Content-Transfer-Encoding"), + "filename": part_data.get_filename(), + "part": part_key, + "uid": uid, + } + + +@callback +def async_setup_services(hass: HomeAssistant) -> None: + """Register the imap services.""" + hass.services.async_register(DOMAIN, "seen", _async_seen, SERVICE_SEEN_SCHEMA) + hass.services.async_register(DOMAIN, "move", _async_move, SERVICE_MOVE_SCHEMA) + hass.services.async_register(DOMAIN, "delete", _async_delete, SERVICE_DELETE_SCHEMA) + hass.services.async_register( + DOMAIN, + "fetch", + _async_fetch, + SERVICE_FETCH_TEXT_SCHEMA, + supports_response=SupportsResponse.ONLY, + ) + hass.services.async_register( + DOMAIN, + "fetch_part", + _async_fetch_part, + SERVICE_FETCH_PART_SCHEMA, + supports_response=SupportsResponse.ONLY, + ) diff --git a/tests/components/imap/test_init.py b/tests/components/imap/test_init.py index e7628b72c073..27f8c52d56a8 100644 --- a/tests/components/imap/test_init.py +++ b/tests/components/imap/test_init.py @@ -1034,7 +1034,8 @@ async def test_services( } for translation_key, attrs in exceptions.items(): with patch( - "homeassistant.components.imap.connect_to_server", side_effect=attrs["exc"] + "homeassistant.components.imap.services.connect_to_server", + side_effect=attrs["exc"], ): data = {"entry": config_entry.entry_id, "uid": "1"} with pytest.raises(ServiceValidationError) as exc: