From de6c3512d2c4cab88940da98a8e3d3331ddbf838 Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Sat, 27 Sep 2025 14:49:26 +0200 Subject: [PATCH] Add IMAP fetch message part feature (#152845) --- homeassistant/components/imap/__init__.py | 78 ++++++++++ homeassistant/components/imap/coordinator.py | 25 +++- homeassistant/components/imap/icons.json | 3 + homeassistant/components/imap/services.yaml | 19 +++ homeassistant/components/imap/strings.json | 21 +++ tests/components/imap/const.py | 67 +++++++++ tests/components/imap/test_init.py | 149 ++++++++++++++++--- 7 files changed, 342 insertions(+), 20 deletions(-) diff --git a/homeassistant/components/imap/__init__.py b/homeassistant/components/imap/__init__.py index 5349f249ab3..a60bc308410 100644 --- a/homeassistant/components/imap/__init__.py +++ b/homeassistant/components/imap/__init__.py @@ -3,7 +3,9 @@ from __future__ import annotations import asyncio +from email.message import Message import logging +from typing import Any from aioimaplib import IMAP4_SSL, AioImapException, Response import voluptuous as vol @@ -33,6 +35,7 @@ from .coordinator import ( ImapPollingDataUpdateCoordinator, ImapPushDataUpdateCoordinator, connect_to_server, + get_parts, ) from .errors import InvalidAuth, InvalidFolder @@ -40,6 +43,7 @@ PLATFORMS: list[Platform] = [Platform.SENSOR] CONF_ENTRY = "entry" CONF_SEEN = "seen" +CONF_PART = "part" CONF_UID = "uid" CONF_TARGET_FOLDER = "target_folder" @@ -64,6 +68,11 @@ SERVICE_MOVE_SCHEMA = _SERVICE_UID_SCHEMA.extend( ) SERVICE_DELETE_SCHEMA = _SERVICE_UID_SCHEMA SERVICE_FETCH_TEXT_SCHEMA = _SERVICE_UID_SCHEMA +SERVICE_FETCH_PART_SCHEMA = _SERVICE_UID_SCHEMA.extend( + { + vol.Required(CONF_PART): cv.string, + } +) type ImapConfigEntry = ConfigEntry[ImapDataUpdateCoordinator] @@ -216,12 +225,14 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: 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), "uid": uid, } @@ -233,6 +244,73 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: 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, + ) + return True diff --git a/homeassistant/components/imap/coordinator.py b/homeassistant/components/imap/coordinator.py index 34d3f43eb69..af8fcc91155 100644 --- a/homeassistant/components/imap/coordinator.py +++ b/homeassistant/components/imap/coordinator.py @@ -21,7 +21,7 @@ from homeassistant.const import ( CONF_VERIFY_SSL, CONTENT_TYPE_TEXT_PLAIN, ) -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ( ConfigEntryAuthFailed, ConfigEntryError, @@ -209,6 +209,28 @@ class ImapMessage: return str(self.email_message.get_payload()) +@callback +def get_parts(message: Message, prefix: str | None = None) -> dict[str, Any]: + """Return information about the parts of a multipart message.""" + parts: dict[str, Any] = {} + if not message.is_multipart(): + return {} + for index, part in enumerate(message.get_payload(), 0): + if TYPE_CHECKING: + assert isinstance(part, Message) + key = f"{prefix},{index}" if prefix else f"{index}" + if part.is_multipart(): + parts |= get_parts(part, key) + continue + parts[key] = {"content_type": part.get_content_type()} + if filename := part.get_filename(): + parts[key]["filename"] = filename + if content_transfer_encoding := part.get("Content-Transfer-Encoding"): + parts[key]["content_transfer_encoding"] = content_transfer_encoding + + return parts + + class ImapDataUpdateCoordinator(DataUpdateCoordinator[int | None]): """Base class for imap client.""" @@ -275,6 +297,7 @@ class ImapDataUpdateCoordinator(DataUpdateCoordinator[int | None]): "sender": message.sender, "subject": message.subject, "uid": last_message_uid, + "parts": get_parts(message.email_message), } data.update({key: getattr(message, key) for key in self._event_data_keys}) if self.custom_event_template is not None: diff --git a/homeassistant/components/imap/icons.json b/homeassistant/components/imap/icons.json index 17a11d0fe22..5c134b8ef81 100644 --- a/homeassistant/components/imap/icons.json +++ b/homeassistant/components/imap/icons.json @@ -21,6 +21,9 @@ }, "fetch": { "service": "mdi:email-sync-outline" + }, + "fetch_part": { + "service": "mdi:email-sync-outline" } } } diff --git a/homeassistant/components/imap/services.yaml b/homeassistant/components/imap/services.yaml index be56eb148da..7854a6fd688 100644 --- a/homeassistant/components/imap/services.yaml +++ b/homeassistant/components/imap/services.yaml @@ -56,3 +56,22 @@ fetch: example: "12" selector: text: + +fetch_part: + fields: + entry: + required: true + selector: + config_entry: + integration: "imap" + uid: + required: true + example: "12" + selector: + text: + + part: + required: true + example: "0,1" + selector: + text: diff --git a/homeassistant/components/imap/strings.json b/homeassistant/components/imap/strings.json index 0f6f99dff65..417afcf1756 100644 --- a/homeassistant/components/imap/strings.json +++ b/homeassistant/components/imap/strings.json @@ -84,6 +84,9 @@ "imap_server_fail": { "message": "The IMAP server failed to connect: {error}." }, + "invalid_part_index": { + "message": "Invalid part index." + }, "seen_failed": { "message": "Marking message as seen failed with \"{error}\"." } @@ -148,6 +151,24 @@ } } }, + "fetch_part": { + "name": "Fetch message part", + "description": "Fetches a message part or attachment from an email message.", + "fields": { + "entry": { + "name": "[%key:component::imap::services::fetch::fields::entry::name%]", + "description": "[%key:component::imap::services::fetch::fields::entry::description%]" + }, + "uid": { + "name": "[%key:component::imap::services::fetch::fields::uid::name%]", + "description": "[%key:component::imap::services::fetch::fields::uid::description%]" + }, + "part": { + "name": "Part", + "description": "The message part index." + } + } + }, "seen": { "name": "Mark message as seen", "description": "Marks an email as seen.", diff --git a/tests/components/imap/const.py b/tests/components/imap/const.py index 8f6761bd795..5ddf86153cb 100644 --- a/tests/components/imap/const.py +++ b/tests/components/imap/const.py @@ -27,6 +27,9 @@ TEST_MESSAGE_HEADERS2 = ( TEST_MULTIPART_HEADER = ( b'Content-Type: multipart/related;\r\n\tboundary="Mark=_100584970350292485166"' ) +TEST_MULTIPART_ATTACHMENT_HEADER = ( + b'Content-Type: multipart/mixed; boundary="------------qIuh0xG6dsImymfJo6f2M4Zv"' +) TEST_MESSAGE_HEADERS3 = b"" @@ -36,6 +39,13 @@ TEST_MESSAGE_MULTIPART = ( TEST_MESSAGE_HEADERS1 + DATE_HEADER1 + TEST_MESSAGE_HEADERS2 + TEST_MULTIPART_HEADER ) +TEST_MESSAGE_MULTIPART_ATTACHMENT = ( + TEST_MESSAGE_HEADERS1 + + DATE_HEADER1 + + TEST_MESSAGE_HEADERS2 + + TEST_MULTIPART_ATTACHMENT_HEADER +) + TEST_MESSAGE_NO_SUBJECT_TO_FROM = ( TEST_MESSAGE_HEADERS1 + DATE_HEADER1 + TEST_MESSAGE_HEADERS3 ) @@ -140,6 +150,45 @@ TEST_CONTENT_MULTIPART_BASE64_INVALID = ( + b"\r\n--Mark=_100584970350292485166--\r\n" ) +TEST_CONTENT_MULTIPART_WITH_ATTACHMENT = b""" +\nThis is a multi-part message in MIME format. +--------------qIuh0xG6dsImymfJo6f2M4Zv +Content-Type: multipart/alternative; + boundary="------------N4zNjp2QWnOfrYQhtLL02Bk1" + +--------------N4zNjp2QWnOfrYQhtLL02Bk1 +Content-Type: text/plain; charset=UTF-8; format=flowed +Content-Transfer-Encoding: 7bit + +*Multi* part Test body + +--------------N4zNjp2QWnOfrYQhtLL02Bk1 +Content-Type: text/html; charset=UTF-8 +Content-Transfer-Encoding: 7bit + + + + + + + + +

Multi part Test body

+ + + +--------------N4zNjp2QWnOfrYQhtLL02Bk1-- +--------------qIuh0xG6dsImymfJo6f2M4Zv +Content-Type: text/plain; charset=UTF-8; name="Text attachment content.txt" +Content-Disposition: attachment; filename="Text attachment content.txt" +Content-Transfer-Encoding: base64 + +VGV4dCBhdHRhY2htZW50IGNvbnRlbnQ= + +--------------qIuh0xG6dsImymfJo6f2M4Zv-- +""" + + EMPTY_SEARCH_RESPONSE = ("OK", [b"", b"Search completed (0.0001 + 0.000 secs)."]) EMPTY_SEARCH_RESPONSE_ALT = ("OK", [b"Search completed (0.0001 + 0.000 secs)."]) @@ -303,6 +352,24 @@ TEST_FETCH_RESPONSE_MULTIPART_BASE64 = ( b"Fetch completed (0.0001 + 0.000 secs).", ], ) +TEST_FETCH_RESPONSE_MULTIPART_WITH_ATTACHMENT = ( + "OK", + [ + b"1 FETCH (BODY[] {" + + str( + len( + TEST_MESSAGE_MULTIPART_ATTACHMENT + + TEST_CONTENT_MULTIPART_WITH_ATTACHMENT + ) + ).encode("utf-8") + + b"}", + bytearray( + TEST_MESSAGE_MULTIPART_ATTACHMENT + TEST_CONTENT_MULTIPART_WITH_ATTACHMENT + ), + b")", + b"Fetch completed (0.0001 + 0.000 secs).", + ], +) TEST_FETCH_RESPONSE_MULTIPART_BASE64_INVALID = ( "OK", diff --git a/tests/components/imap/test_init.py b/tests/components/imap/test_init.py index bdd29f7442b..dc5727991c1 100644 --- a/tests/components/imap/test_init.py +++ b/tests/components/imap/test_init.py @@ -1,6 +1,7 @@ """Test the imap entry initialization.""" import asyncio +from base64 import b64decode from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock, call, patch @@ -31,6 +32,7 @@ from .const import ( TEST_FETCH_RESPONSE_MULTIPART_BASE64, TEST_FETCH_RESPONSE_MULTIPART_BASE64_INVALID, TEST_FETCH_RESPONSE_MULTIPART_EMPTY_PLAIN, + TEST_FETCH_RESPONSE_MULTIPART_WITH_ATTACHMENT, TEST_FETCH_RESPONSE_NO_SUBJECT_TO_FROM, TEST_FETCH_RESPONSE_TEXT_BARE, TEST_FETCH_RESPONSE_TEXT_OTHER, @@ -107,20 +109,72 @@ async def test_entry_startup_fails( @pytest.mark.parametrize("imap_search", [TEST_SEARCH_RESPONSE]) @pytest.mark.parametrize( - ("imap_fetch", "valid_date"), + ("imap_fetch", "valid_date", "parts"), [ - (TEST_FETCH_RESPONSE_TEXT_BARE, True), - (TEST_FETCH_RESPONSE_TEXT_PLAIN, True), - (TEST_FETCH_RESPONSE_TEXT_PLAIN_ALT, True), - (TEST_FETCH_RESPONSE_INVALID_DATE1, False), - (TEST_FETCH_RESPONSE_INVALID_DATE2, False), - (TEST_FETCH_RESPONSE_INVALID_DATE3, False), - (TEST_FETCH_RESPONSE_TEXT_OTHER, True), - (TEST_FETCH_RESPONSE_HTML, True), - (TEST_FETCH_RESPONSE_MULTIPART, True), - (TEST_FETCH_RESPONSE_MULTIPART_EMPTY_PLAIN, True), - (TEST_FETCH_RESPONSE_MULTIPART_BASE64, True), - (TEST_FETCH_RESPONSE_BINARY, True), + (TEST_FETCH_RESPONSE_TEXT_BARE, True, {}), + (TEST_FETCH_RESPONSE_TEXT_PLAIN, True, {}), + (TEST_FETCH_RESPONSE_TEXT_PLAIN_ALT, True, {}), + (TEST_FETCH_RESPONSE_INVALID_DATE1, False, {}), + (TEST_FETCH_RESPONSE_INVALID_DATE2, False, {}), + (TEST_FETCH_RESPONSE_INVALID_DATE3, False, {}), + (TEST_FETCH_RESPONSE_TEXT_OTHER, True, {}), + (TEST_FETCH_RESPONSE_HTML, True, {}), + ( + TEST_FETCH_RESPONSE_MULTIPART, + True, + { + "0": { + "content_type": "text/plain", + "content_transfer_encoding": "7bit", + }, + "1": {"content_type": "text/html", "content_transfer_encoding": "7bit"}, + }, + ), + ( + TEST_FETCH_RESPONSE_MULTIPART_EMPTY_PLAIN, + True, + { + "0": { + "content_type": "text/plain", + "content_transfer_encoding": "7bit", + }, + "1": {"content_type": "text/html", "content_transfer_encoding": "7bit"}, + }, + ), + ( + TEST_FETCH_RESPONSE_MULTIPART_BASE64, + True, + { + "0": { + "content_type": "text/plain", + "content_transfer_encoding": "base64", + }, + "1": { + "content_type": "text/html", + "content_transfer_encoding": "base64", + }, + }, + ), + ( + TEST_FETCH_RESPONSE_MULTIPART_WITH_ATTACHMENT, + True, + { + "0,0": { + "content_type": "text/plain", + "content_transfer_encoding": "7bit", + }, + "0,1": { + "content_type": "text/html", + "content_transfer_encoding": "7bit", + }, + "1": { + "content_type": "text/plain", + "filename": "Text attachment content.txt", + "content_transfer_encoding": "base64", + }, + }, + ), + (TEST_FETCH_RESPONSE_BINARY, True, {}), ], ids=[ "bare", @@ -134,13 +188,18 @@ async def test_entry_startup_fails( "multipart", "multipart_empty_plain", "multipart_base64", + "multipart_attachment", "binary", ], ) @pytest.mark.parametrize("imap_has_capability", [True, False], ids=["push", "poll"]) @pytest.mark.parametrize("charset", ["utf-8", "us-ascii"], ids=["utf-8", "us-ascii"]) async def test_receiving_message_successfully( - hass: HomeAssistant, mock_imap_protocol: MagicMock, valid_date: bool, charset: str + hass: HomeAssistant, + mock_imap_protocol: MagicMock, + valid_date: bool, + charset: str, + parts: dict[str, Any], ) -> None: """Test receiving a message successfully.""" event_called = async_capture_events(hass, "imap_content") @@ -170,6 +229,7 @@ async def test_receiving_message_successfully( assert data["sender"] == "john.doe@example.com" assert data["subject"] == "Test subject" assert data["uid"] == "1" + assert data["parts"] == parts assert "Test body" in data["text"] assert (valid_date and isinstance(data["date"], datetime)) or ( not valid_date and data["date"] is None @@ -826,11 +886,33 @@ async def test_enforce_polling( @pytest.mark.parametrize( - ("imap_search", "imap_fetch"), - [(TEST_SEARCH_RESPONSE, TEST_FETCH_RESPONSE_TEXT_PLAIN)], + ("imap_search", "imap_fetch", "message_parts"), + [ + ( + TEST_SEARCH_RESPONSE, + TEST_FETCH_RESPONSE_MULTIPART_WITH_ATTACHMENT, + { + "0,0": { + "content_type": "text/plain", + "content_transfer_encoding": "7bit", + }, + "0,1": { + "content_type": "text/html", + "content_transfer_encoding": "7bit", + }, + "1": { + "content_type": "text/plain", + "filename": "Text attachment content.txt", + "content_transfer_encoding": "base64", + }, + }, + ) + ], ) @pytest.mark.parametrize("imap_has_capability", [True, False], ids=["push", "poll"]) -async def test_services(hass: HomeAssistant, mock_imap_protocol: MagicMock) -> None: +async def test_services( + hass: HomeAssistant, mock_imap_protocol: MagicMock, message_parts: dict[str, Any] +) -> None: """Test receiving a message successfully.""" event_called = async_capture_events(hass, "imap_content") @@ -859,6 +941,7 @@ async def test_services(hass: HomeAssistant, mock_imap_protocol: MagicMock) -> N assert data["subject"] == "Test subject" assert data["uid"] == "1" assert data["entry_id"] == config_entry.entry_id + assert data["parts"] == message_parts # Test seen service data = {"entry": config_entry.entry_id, "uid": "1"} @@ -889,16 +972,42 @@ async def test_services(hass: HomeAssistant, mock_imap_protocol: MagicMock) -> N mock_imap_protocol.store.assert_called_with("1", "+FLAGS (\\Deleted)") mock_imap_protocol.protocol.expunge.assert_called_once() - # Test fetch service + # Test fetch service with text response + mock_imap_protocol.reset_mock() data = {"entry": config_entry.entry_id, "uid": "1"} response = await hass.services.async_call( DOMAIN, "fetch", data, blocking=True, return_response=True ) mock_imap_protocol.fetch.assert_called_with("1", "BODY.PEEK[]") - assert response["text"] == "Test body\r\n" + assert response["text"] == "*Multi* part Test body\n" assert response["sender"] == "john.doe@example.com" assert response["subject"] == "Test subject" assert response["uid"] == "1" + assert response["parts"] == message_parts + + # Test fetch part service with attachment response + mock_imap_protocol.reset_mock() + data = {"entry": config_entry.entry_id, "uid": "1", "part": "1"} + response = await hass.services.async_call( + DOMAIN, "fetch_part", data, blocking=True, return_response=True + ) + mock_imap_protocol.fetch.assert_called_with("1", "BODY.PEEK[]") + assert response["part_data"] == "VGV4dCBhdHRhY2htZW50IGNvbnRlbnQ=\n" + assert response["content_type"] == "text/plain" + assert response["content_transfer_encoding"] == "base64" + assert response["filename"] == "Text attachment content.txt" + assert response["part"] == "1" + assert response["uid"] == "1" + assert b64decode(response["part_data"]) == b"Text attachment content" + + # Test fetch part service with invalid part index + for part in ("A", "2", "0"): + data = {"entry": config_entry.entry_id, "uid": "1", "part": part} + with pytest.raises(ServiceValidationError) as exc: + await hass.services.async_call( + DOMAIN, "fetch_part", data, blocking=True, return_response=True + ) + assert exc.value.translation_key == "invalid_part_index" # Test with invalid entry_id data = {"entry": "invalid", "uid": "1"} @@ -943,12 +1052,14 @@ async def test_services(hass: HomeAssistant, mock_imap_protocol: MagicMock) -> N ), "delete": ({"entry": config_entry.entry_id, "uid": "1"}, False), "fetch": ({"entry": config_entry.entry_id, "uid": "1"}, True), + "fetch_part": ({"entry": config_entry.entry_id, "uid": "1", "part": "1"}, True), } patch_error_translation_key = { "seen": ("store", "seen_failed"), "move": ("copy", "copy_failed"), "delete": ("store", "delete_failed"), "fetch": ("fetch", "fetch_failed"), + "fetch_part": ("fetch", "fetch_failed"), } for service, (data, response) in service_calls_response.items(): with (