From 183c191d63d945ddf7b7d0a442696bf81081b413 Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Fri, 16 Aug 2024 13:34:14 +0200 Subject: [PATCH] Allow raw mqtt payload to be in mqtt publish action (#123900) * Publish raw rendered mqtt payload as raw for mqtt publish action * Move check out of try block * Only try to eval `bytes` is payload starts with supported string Co-authored-by: Erik Montnemery * Improve docst * Add `evaluate_bytes` option to publish action * Rename to `evaluate_payload` * Update homeassistant/components/mqtt/strings.json Co-authored-by: Erik Montnemery * Extend test to assert literal_eval is called or not --------- Co-authored-by: Erik Montnemery --- homeassistant/components/mqtt/__init__.py | 7 +++ homeassistant/components/mqtt/models.py | 34 +++++------ homeassistant/components/mqtt/services.yaml | 5 ++ homeassistant/components/mqtt/strings.json | 4 ++ tests/components/mqtt/test_init.py | 68 +++++++++++++++++++++ 5 files changed, 101 insertions(+), 17 deletions(-) diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py index 013bd26e49c..b2adb7665fc 100644 --- a/homeassistant/components/mqtt/__init__.py +++ b/homeassistant/components/mqtt/__init__.py @@ -90,6 +90,7 @@ from .models import ( # noqa: F401 PublishPayloadType, ReceiveMessage, ReceivePayloadType, + convert_outgoing_mqtt_payload, ) from .subscription import ( # noqa: F401 EntitySubscription, @@ -115,6 +116,7 @@ SERVICE_DUMP = "dump" ATTR_TOPIC_TEMPLATE = "topic_template" ATTR_PAYLOAD_TEMPLATE = "payload_template" +ATTR_EVALUATE_PAYLOAD = "evaluate_payload" MAX_RECONNECT_WAIT = 300 # seconds @@ -166,6 +168,7 @@ MQTT_PUBLISH_SCHEMA = vol.All( vol.Exclusive(ATTR_TOPIC_TEMPLATE, CONF_TOPIC): cv.string, vol.Exclusive(ATTR_PAYLOAD, CONF_PAYLOAD): cv.string, vol.Exclusive(ATTR_PAYLOAD_TEMPLATE, CONF_PAYLOAD): cv.string, + vol.Optional(ATTR_EVALUATE_PAYLOAD): cv.boolean, vol.Optional(ATTR_QOS, default=DEFAULT_QOS): valid_qos_schema, vol.Optional(ATTR_RETAIN, default=DEFAULT_RETAIN): cv.boolean, }, @@ -295,6 +298,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: msg_topic: str | None = call.data.get(ATTR_TOPIC) msg_topic_template: str | None = call.data.get(ATTR_TOPIC_TEMPLATE) payload: PublishPayloadType = call.data.get(ATTR_PAYLOAD) + evaluate_payload: bool = call.data.get(ATTR_EVALUATE_PAYLOAD, False) payload_template: str | None = call.data.get(ATTR_PAYLOAD_TEMPLATE) qos: int = call.data[ATTR_QOS] retain: bool = call.data[ATTR_RETAIN] @@ -354,6 +358,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: payload = MqttCommandTemplate( template.Template(payload_template, hass) ).async_render() + elif evaluate_payload: + # Convert quoted binary literal to raw data + payload = convert_outgoing_mqtt_payload(payload) if TYPE_CHECKING: assert msg_topic is not None diff --git a/homeassistant/components/mqtt/models.py b/homeassistant/components/mqtt/models.py index f2b3165f66c..f7abbc29464 100644 --- a/homeassistant/components/mqtt/models.py +++ b/homeassistant/components/mqtt/models.py @@ -51,6 +51,22 @@ ATTR_THIS = "this" type PublishPayloadType = str | bytes | int | float | None +def convert_outgoing_mqtt_payload( + payload: PublishPayloadType, +) -> PublishPayloadType: + """Ensure correct raw MQTT payload is passed as bytes for publishing.""" + if isinstance(payload, str) and payload.startswith(("b'", 'b"')): + try: + native_object = literal_eval(payload) + except (ValueError, TypeError, SyntaxError, MemoryError): + pass + else: + if isinstance(native_object, bytes): + return native_object + + return payload + + @dataclass class PublishMessage: """MQTT Message for publishing.""" @@ -173,22 +189,6 @@ class MqttCommandTemplate: variables: TemplateVarsType = None, ) -> PublishPayloadType: """Render or convert the command template with given value or variables.""" - - def _convert_outgoing_payload( - payload: PublishPayloadType, - ) -> PublishPayloadType: - """Ensure correct raw MQTT payload is passed as bytes for publishing.""" - if isinstance(payload, str): - try: - native_object = literal_eval(payload) - if isinstance(native_object, bytes): - return native_object - - except (ValueError, TypeError, SyntaxError, MemoryError): - pass - - return payload - if self._command_template is None: return value @@ -210,7 +210,7 @@ class MqttCommandTemplate: self._command_template, ) try: - return _convert_outgoing_payload( + return convert_outgoing_mqtt_payload( self._command_template.async_render(values, parse_result=False) ) except TemplateError as exc: diff --git a/homeassistant/components/mqtt/services.yaml b/homeassistant/components/mqtt/services.yaml index ee5e4ff56e8..c5e4f372bd6 100644 --- a/homeassistant/components/mqtt/services.yaml +++ b/homeassistant/components/mqtt/services.yaml @@ -12,6 +12,11 @@ publish: example: "The temperature is {{ states('sensor.temperature') }}" selector: template: + evaluate_payload: + advanced: true + default: false + selector: + boolean: qos: advanced: true default: 0 diff --git a/homeassistant/components/mqtt/strings.json b/homeassistant/components/mqtt/strings.json index 93131376154..c786d7e08a1 100644 --- a/homeassistant/components/mqtt/strings.json +++ b/homeassistant/components/mqtt/strings.json @@ -230,6 +230,10 @@ "name": "Publish", "description": "Publishes a message to an MQTT topic.", "fields": { + "evaluate_payload": { + "name": "Evaluate payload", + "description": "When `payload` is a Python bytes literal, evaluate the bytes literal and publish the raw data." + }, "topic": { "name": "Topic", "description": "Topic to publish to." diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py index f495c5ca585..333960d8ad4 100644 --- a/tests/components/mqtt/test_init.py +++ b/tests/components/mqtt/test_init.py @@ -420,6 +420,74 @@ async def test_mqtt_publish_action_call_with_template_payload_renders_template( mqtt_mock.reset_mock() +@pytest.mark.parametrize( + ("attr_payload", "payload", "evaluate_payload", "literal_eval_calls"), + [ + ("b'\\xde\\xad\\xbe\\xef'", b"\xde\xad\xbe\xef", True, 1), + ("b'\\xde\\xad\\xbe\\xef'", "b'\\xde\\xad\\xbe\\xef'", False, 0), + ("DEADBEEF", "DEADBEEF", False, 0), + ( + "b'\\xde", + "b'\\xde", + True, + 1, + ), # Bytes literal is invalid, fall back to string + ], +) +async def test_mqtt_publish_action_call_with_raw_data( + hass: HomeAssistant, + mqtt_mock_entry: MqttMockHAClientGenerator, + attr_payload: str, + payload: str | bytes, + evaluate_payload: bool, + literal_eval_calls: int, +) -> None: + """Test the mqtt publish action call raw data. + + When `payload` represents a `bytes` object, it should be published + as raw data if `evaluate_payload` is set. + """ + mqtt_mock = await mqtt_mock_entry() + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC: "test/topic", + mqtt.ATTR_PAYLOAD: attr_payload, + mqtt.ATTR_EVALUATE_PAYLOAD: evaluate_payload, + }, + blocking=True, + ) + assert mqtt_mock.async_publish.called + assert mqtt_mock.async_publish.call_args[0][1] == payload + + with patch( + "homeassistant.components.mqtt.models.literal_eval" + ) as literal_eval_mock: + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC: "test/topic", + mqtt.ATTR_PAYLOAD: attr_payload, + }, + blocking=True, + ) + literal_eval_mock.assert_not_called() + + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC: "test/topic", + mqtt.ATTR_PAYLOAD: attr_payload, + mqtt.ATTR_EVALUATE_PAYLOAD: evaluate_payload, + }, + blocking=True, + ) + assert len(literal_eval_mock.mock_calls) == literal_eval_calls + + # The use of a payload_template in an mqtt publish action call # has been deprecated with HA Core 2024.8.0 and will be removed with HA Core 2025.2.0 async def test_publish_action_call_with_bad_payload_template(