1
0
mirror of https://github.com/home-assistant/core.git synced 2026-02-15 07:36:16 +00:00

Add IMAP fetch message part feature (#152845)

This commit is contained in:
Jan Bouwhuis
2025-09-27 14:49:26 +02:00
committed by GitHub
parent 36dc1e938a
commit de6c3512d2
7 changed files with 342 additions and 20 deletions

View File

@@ -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

View File

@@ -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:

View File

@@ -21,6 +21,9 @@
},
"fetch": {
"service": "mdi:email-sync-outline"
},
"fetch_part": {
"service": "mdi:email-sync-outline"
}
}
}

View File

@@ -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:

View File

@@ -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.",

View File

@@ -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
<!DOCTYPE html>
<html>
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8">
</head>
<body>
<p><b>Multi</b> part Test body</p>
</body>
</html>
--------------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",

View File

@@ -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 (