Add advanced imap option to set custom event max message size (#93163)

homeassistant/components/imap/config_flow.py

@@ -24,11 +24,14 @@ from homeassistant.util.ssl import SSLCipherList
 from .const import (
     CONF_CHARSET,
     CONF_FOLDER,
+    CONF_MAX_MESSAGE_SIZE,
     CONF_SEARCH,
     CONF_SERVER,
     CONF_SSL_CIPHER_LIST,
+    DEFAULT_MAX_MESSAGE_SIZE,
     DEFAULT_PORT,
     DOMAIN,
+    MAX_MESSAGE_SIZE_LIMIT,
 )
 from .coordinator import connect_to_server
 from .errors import InvalidAuth, InvalidFolder
@@ -55,7 +58,7 @@ CONFIG_SCHEMA = vol.Schema(
 CONFIG_SCHEMA_ADVANCED = {
     vol.Optional(
         CONF_SSL_CIPHER_LIST, default=SSLCipherList.PYTHON_DEFAULT
-    ): CIPHER_SELECTOR
+    ): CIPHER_SELECTOR,
 }
 
 OPTIONS_SCHEMA = vol.Schema(
@@ -65,6 +68,13 @@ OPTIONS_SCHEMA = vol.Schema(
     }
 )
 
+OPTIONS_SCHEMA_ADVANCED = {
+    vol.Optional(CONF_MAX_MESSAGE_SIZE, default=DEFAULT_MAX_MESSAGE_SIZE): vol.All(
+        cv.positive_int,
+        vol.Range(min=DEFAULT_MAX_MESSAGE_SIZE, max=MAX_MESSAGE_SIZE_LIMIT),
+    )
+}
+
 
 async def validate_input(user_input: dict[str, Any]) -> dict[str, str]:
     """Validate user input."""
@@ -233,6 +243,9 @@ class OptionsFlow(config_entries.OptionsFlowWithConfigEntry):
                     )
                     return self.async_create_entry(data={})
 
-        schema = self.add_suggested_values_to_schema(OPTIONS_SCHEMA, entry_data)
+        schema = OPTIONS_SCHEMA
+        if self.show_advanced_options:
+            schema = schema.extend(OPTIONS_SCHEMA_ADVANCED)
+        schema = self.add_suggested_values_to_schema(schema, entry_data)
 
         return self.async_show_form(step_id="init", data_schema=schema, errors=errors)

homeassistant/components/imap/const.py

@@ -8,6 +8,11 @@ CONF_SERVER: Final = "server"
 CONF_FOLDER: Final = "folder"
 CONF_SEARCH: Final = "search"
 CONF_CHARSET: Final = "charset"
+CONF_MAX_MESSAGE_SIZE = "max_message_size"
 CONF_SSL_CIPHER_LIST: Final = "ssl_cipher_list"
 
 DEFAULT_PORT: Final = 993
+
+DEFAULT_MAX_MESSAGE_SIZE = 2048
+
+MAX_MESSAGE_SIZE_LIMIT = 30000

homeassistant/components/imap/coordinator.py

@@ -20,15 +20,18 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryError
+from homeassistant.helpers.json import json_bytes
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from homeassistant.util.ssl import SSLCipherList, client_context
 
 from .const import (
     CONF_CHARSET,
     CONF_FOLDER,
+    CONF_MAX_MESSAGE_SIZE,
     CONF_SEARCH,
     CONF_SERVER,
     CONF_SSL_CIPHER_LIST,
+    DEFAULT_MAX_MESSAGE_SIZE,
     DOMAIN,
 )
 from .errors import InvalidAuth, InvalidFolder
@@ -38,6 +41,7 @@ _LOGGER = logging.getLogger(__name__)
 BACKOFF_TIME = 10
 
 EVENT_IMAP = "imap_content"
+MAX_EVENT_DATA_BYTES = 32168
 
 
 async def connect_to_server(data: Mapping[str, Any]) -> IMAP4_SSL:
@@ -177,11 +181,26 @@ class ImapDataUpdateCoordinator(DataUpdateCoordinator[int | None]):
                 "search": self.config_entry.data[CONF_SEARCH],
                 "folder": self.config_entry.data[CONF_FOLDER],
                 "date": message.date,
-                "text": message.text[:2048],
+                "text": message.text[
+                    : self.config_entry.data.get(
+                        CONF_MAX_MESSAGE_SIZE, DEFAULT_MAX_MESSAGE_SIZE
+                    )
+                ],
                 "sender": message.sender,
                 "subject": message.subject,
                 "headers": message.headers,
             }
+            if (size := len(json_bytes(data))) > MAX_EVENT_DATA_BYTES:
+                _LOGGER.warning(
+                    "Custom imap_content event skipped, size (%s) exceeds "
+                    "the maximal event size (%s), sender: %s, subject: %s",
+                    size,
+                    MAX_EVENT_DATA_BYTES,
+                    message.sender,
+                    message.subject,
+                )
+                return
+
             self.hass.bus.fire(EVENT_IMAP, data)
             _LOGGER.debug(
                 "Message processed, sender: %s, subject: %s",

homeassistant/components/imap/strings.json

@@ -39,7 +39,8 @@
       "init": {
         "data": {
           "folder": "[%key:component::imap::config::step::user::data::folder%]",
-          "search": "[%key:component::imap::config::step::user::data::search%]"
+          "search": "[%key:component::imap::config::step::user::data::search%]",
+          "max_message_size": "Max message size (2048 < size < 30000)"
         }
       }
     },

tests/components/imap/test_config_flow.py

@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, patch
 
 from aioimaplib import AioImapException
 import pytest
+import voluptuous as vol
 
 from homeassistant import config_entries, data_entry_flow
 from homeassistant.components.imap.const import (
@@ -397,6 +398,54 @@ async def test_key_options_in_options_form(hass: HomeAssistant) -> None:
     assert result2["errors"] == {"base": "already_configured"}
 
 
+@pytest.mark.parametrize(
+    ("advanced_options", "assert_result"),
+    [
+        ({"max_message_size": "8192"}, data_entry_flow.FlowResultType.CREATE_ENTRY),
+        ({"max_message_size": "1024"}, data_entry_flow.FlowResultType.FORM),
+        ({"max_message_size": "65536"}, data_entry_flow.FlowResultType.FORM),
+    ],
+)
+async def test_advanced_options_form(
+    hass: HomeAssistant,
+    advanced_options: dict[str, str],
+    assert_result: data_entry_flow.FlowResultType,
+) -> None:
+    """Test we show the advanced options."""
+
+    entry = MockConfigEntry(domain=DOMAIN, data=MOCK_CONFIG)
+    entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(entry.entry_id)
+
+    result = await hass.config_entries.options.async_init(
+        entry.entry_id,
+        context={"source": config_entries.SOURCE_USER, "show_advanced_options": True},
+    )
+
+    assert result["type"] == data_entry_flow.FlowResultType.FORM
+    assert result["step_id"] == "init"
+
+    new_config = MOCK_OPTIONS.copy()
+    new_config.update(advanced_options)
+
+    try:
+        with patch(
+            "homeassistant.components.imap.config_flow.connect_to_server"
+        ) as mock_client:
+            mock_client.return_value.search.return_value = ("OK", [b""])
+            # Option update should fail if FlowResultType.FORM is expected
+            result2 = await hass.config_entries.options.async_configure(
+                result["flow_id"], new_config
+            )
+            assert result2["type"] == assert_result
+            # Check if entry was updated
+            for key, value in new_config.items():
+                assert str(entry.data[key]) == value
+    except vol.MultipleInvalid:
+        # Check if form was expected with these options
+        assert assert_result == data_entry_flow.FlowResultType.FORM
+
+
 async def test_import_flow_success(hass: HomeAssistant) -> None:
     """Test a successful import of yaml."""
     with patch(

tests/components/imap/test_init.py

@@ -446,3 +446,62 @@ async def test_reset_last_message(
 
     # One new event
     assert len(event_called) == 2
+
+
+@pytest.mark.parametrize("imap_search", [TEST_SEARCH_RESPONSE])
+@pytest.mark.parametrize(
+    "imap_fetch", [(TEST_FETCH_RESPONSE_TEXT_PLAIN)], ids=["plain"]
+)
+@pytest.mark.parametrize("imap_has_capability", [True, False], ids=["push", "poll"])
+@patch("homeassistant.components.imap.coordinator.MAX_EVENT_DATA_BYTES", 500)
+async def test_event_skipped_message_too_large(
+    hass: HomeAssistant, mock_imap_protocol: MagicMock, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test skipping event when message is to large."""
+    event_called = async_capture_events(hass, "imap_content")
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data=MOCK_CONFIG)
+    config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+    # Make sure we have had one update (when polling)
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=5))
+    await hass.async_block_till_done()
+    state = hass.states.get("sensor.imap_email_email_com")
+    # We should have received one message
+    assert state is not None
+    assert state.state == "1"
+    assert len(event_called) == 0
+    assert "Custom imap_content event skipped" in caplog.text
+
+
+@pytest.mark.parametrize("imap_search", [TEST_SEARCH_RESPONSE])
+@pytest.mark.parametrize(
+    "imap_fetch", [(TEST_FETCH_RESPONSE_TEXT_PLAIN)], ids=["plain"]
+)
+@pytest.mark.parametrize("imap_has_capability", [True, False], ids=["push", "poll"])
+async def test_message_is_truncated(
+    hass: HomeAssistant, mock_imap_protocol: MagicMock, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test truncating message text in event data."""
+    event_called = async_capture_events(hass, "imap_content")
+
+    config = MOCK_CONFIG.copy()
+
+    # Mock the max message size to test it is truncated
+    config["max_message_size"] = 3
+    config_entry = MockConfigEntry(domain=DOMAIN, data=config)
+    config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+    # Make sure we have had one update (when polling)
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=5))
+    await hass.async_block_till_done()
+    state = hass.states.get("sensor.imap_email_email_com")
+    # We should have received one message
+    assert state is not None
+    assert state.state == "1"
+    assert len(event_called) == 1
+
+    event_data = event_called[0].data
+    assert len(event_data["text"]) == 3