From 5bc825a8ab396a975deac39230af957e19db070c Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Mon, 22 May 2023 12:14:06 +0200 Subject: [PATCH] Add advanced imap option to set custom event max message size (#93163) --- homeassistant/components/imap/config_flow.py | 17 +++++- homeassistant/components/imap/const.py | 5 ++ homeassistant/components/imap/coordinator.py | 21 ++++++- homeassistant/components/imap/strings.json | 3 +- tests/components/imap/test_config_flow.py | 49 ++++++++++++++++ tests/components/imap/test_init.py | 59 ++++++++++++++++++++ 6 files changed, 150 insertions(+), 4 deletions(-) diff --git a/homeassistant/components/imap/config_flow.py b/homeassistant/components/imap/config_flow.py index 71b09048e6f..8724dbf97c0 100644 --- a/homeassistant/components/imap/config_flow.py +++ b/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) diff --git a/homeassistant/components/imap/const.py b/homeassistant/components/imap/const.py index a1ca586b48b..b39c8808633 100644 --- a/homeassistant/components/imap/const.py +++ b/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 diff --git a/homeassistant/components/imap/coordinator.py b/homeassistant/components/imap/coordinator.py index 0f53c5928c9..512df9adf51 100644 --- a/homeassistant/components/imap/coordinator.py +++ b/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", diff --git a/homeassistant/components/imap/strings.json b/homeassistant/components/imap/strings.json index e50370dd9b1..6e97fbe69d8 100644 --- a/homeassistant/components/imap/strings.json +++ b/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)" } } }, diff --git a/tests/components/imap/test_config_flow.py b/tests/components/imap/test_config_flow.py index 82430549f05..b8242a87be8 100644 --- a/tests/components/imap/test_config_flow.py +++ b/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( diff --git a/tests/components/imap/test_init.py b/tests/components/imap/test_init.py index 8f00cf395d2..26c0325a50c 100644 --- a/tests/components/imap/test_init.py +++ b/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