prepare-for-shutdown: publish config for home assistant to enable aut…

…omatic discovery

CHANGELOG.md

@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- home assistant: enable [automatic discovery](https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix)
+  for logind's `PreparingForShutdown` signal
 
 ## [0.2.0] - 2020-06-21
 ### Added

README.md

@@ -71,6 +71,24 @@ automation:
     entity_id: switch.desk_lamp
 ```
 
+### Automatic Discovery of Shutdown Sensor (Optional)
+
+After enabling [MQTT device discovery](https://www.home-assistant.io/docs/mqtt/discovery/)
+home assistant will automatically detect a new entity
+`binary_sensor.hostname_preparing_for_shutdown`.
+
+```yaml
+mqtt:
+  broker: BROKER_HOSTNAME_OR_IP_ADDRESS
+  discovery: true
+  # credentials, additional options…
+```
+
+![homeassistant discovery binary_sensor.hostname_preparing_for_shutdown](docs/homeassistant/preparing-for-shutdown/settings/discovery/2020-06-21.png)
+
+When using a custom `discovery_prefix`
+pass `--homeassistant-discovery-prefix custom-prefix` to `systemctl-mqtt`.
+
 ## Docker 🐳
 
 1. Clone this repository.

systemctl_mqtt/__init__.py

@@ -35,6 +35,9 @@
 import gi.repository.GLib  # pylint-import-requirements: imports=PyGObject
 import paho.mqtt.client
 
+import systemctl_mqtt._homeassistant
+import systemctl_mqtt._mqtt
+
 _LOGGER = logging.getLogger(__name__)
 
 _SHUTDOWN_DELAY = datetime.timedelta(seconds=4)
@@ -115,8 +118,15 @@ def _schedule_shutdown(action: str) -> None:
 
 
 class _State:
-    def __init__(self, mqtt_topic_prefix: str) -> None:
+    def __init__(
+        self,
+        mqtt_topic_prefix: str,
+        homeassistant_discovery_prefix: str,
+        homeassistant_node_id: str,
+    ) -> None:
         self._mqtt_topic_prefix = mqtt_topic_prefix
+        self._homeassistant_discovery_prefix = homeassistant_discovery_prefix
+        self._homeassistant_node_id = homeassistant_node_id
         self._login_manager = _get_login_manager()  # type: dbus.proxies.Interface
         self._shutdown_lock = None  # type: typing.Optional[dbus.types.UnixFd]
         self._shutdown_lock_mutex = threading.Lock()
@@ -142,12 +152,17 @@ def release_shutdown_lock(self) -> None:
                 _LOGGER.debug("released shutdown inhibitor lock")
                 self._shutdown_lock = None
 
+    @property
+    def _preparing_for_shutdown_topic(self) -> str:
+        return self.mqtt_topic_prefix + "/preparing-for-shutdown"
+
     def _publish_preparing_for_shutdown(
         self, mqtt_client: paho.mqtt.client.Client, active: bool, block: bool,
     ) -> None:
         # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1199
-        topic = self.mqtt_topic_prefix + "/preparing-for-shutdown"
-        payload = json.dumps(active)
+        topic = self._preparing_for_shutdown_topic
+        # pylint: disable=protected-access
+        payload = systemctl_mqtt._mqtt.encode_bool(active)
         _LOGGER.info("publishing %r on %s", payload, topic)
         msg_info = mqtt_client.publish(
             topic=topic, payload=payload, retain=True,
@@ -206,6 +221,43 @@ def publish_preparing_for_shutdown(
             block=False,
         )
 
+    def publish_preparing_for_shutdown_homeassistant_config(
+        self, mqtt_client: paho.mqtt.client.Client
+    ) -> None:
+        # <discovery_prefix>/<component>/[<node_id>/]<object_id>/config
+        # https://www.home-assistant.io/docs/mqtt/discovery/
+        discovery_topic = "/".join(
+            (
+                self._homeassistant_discovery_prefix,
+                "binary_sensor",
+                self._homeassistant_node_id,
+                "preparing-for-shutdown",
+                "config",
+            )
+        )
+        unique_id = "/".join(
+            (
+                "systemctl-mqtt",
+                self._homeassistant_node_id,
+                "logind",
+                "preparing-for-shutdown",
+            )
+        )
+        # https://www.home-assistant.io/integrations/binary_sensor.mqtt/#configuration-variables
+        config = {
+            "unique_id": unique_id,
+            "state_topic": self._preparing_for_shutdown_topic,
+            # pylint: disable=protected-access
+            "payload_on": systemctl_mqtt._mqtt.encode_bool(True),
+            "payload_off": systemctl_mqtt._mqtt.encode_bool(False),
+            # friendly_name & template for default entity_id
+            "name": "{} preparing for shutdown".format(self._homeassistant_node_id),
+        }
+        _LOGGER.debug("publishing home assistant config on %s", discovery_topic)
+        mqtt_client.publish(
+            topic=discovery_topic, payload=json.dumps(config), retain=True,
+        )
+
 
 class _MQTTAction:
 
@@ -254,6 +306,7 @@ def _mqtt_on_connect(
     state.acquire_shutdown_lock()
     state.register_prepare_for_shutdown_handler(mqtt_client=mqtt_client)
     state.publish_preparing_for_shutdown(mqtt_client=mqtt_client)
+    state.publish_preparing_for_shutdown_homeassistant_config(mqtt_client=mqtt_client)
     for topic_suffix, action in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.items():
         topic = state.mqtt_topic_prefix + "/" + topic_suffix
         _LOGGER.info("subscribing to %s", topic)
@@ -272,12 +325,19 @@ def _run(
     mqtt_username: typing.Optional[str],
     mqtt_password: typing.Optional[str],
     mqtt_topic_prefix: str,
+    homeassistant_discovery_prefix: str,
+    homeassistant_node_id: str,
 ) -> None:
+    # pylint: disable=too-many-arguments
     # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
     dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
     # https://pypi.org/project/paho-mqtt/
     mqtt_client = paho.mqtt.client.Client(
-        userdata=_State(mqtt_topic_prefix=mqtt_topic_prefix)
+        userdata=_State(
+            mqtt_topic_prefix=mqtt_topic_prefix,
+            homeassistant_discovery_prefix=homeassistant_discovery_prefix,
+            homeassistant_node_id=homeassistant_node_id,
+        )
     )
     mqtt_client.on_connect = _mqtt_on_connect
     mqtt_client.tls_set(ca_certs=None)  # enable tls trusting default system certs
@@ -303,10 +363,6 @@ def _run(
         _LOGGER.debug("MQTT loop stopped")
 
 
-def _get_hostname() -> str:
-    return socket.gethostname()
-
-
 def _main() -> None:
     logging.basicConfig(
         level=logging.DEBUG,
@@ -329,13 +385,24 @@ def _main() -> None:
         dest="mqtt_password_path",
         help="stripping trailing newline",
     )
-    # https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix
     argparser.add_argument(
         "--mqtt-topic-prefix",
         type=str,
-        default="systemctl/" + _get_hostname(),
+        # pylint: disable=protected-access
+        default="systemctl/" + systemctl_mqtt._utils.get_hostname(),
         help=" ",  # show default
     )
+    # https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix
+    argparser.add_argument(
+        "--homeassistant-discovery-prefix", type=str, default="homeassistant", help=" ",
+    )
+    argparser.add_argument(
+        "--homeassistant-node-id",
+        type=str,
+        # pylint: disable=protected-access
+        default=systemctl_mqtt._homeassistant.get_default_node_id(),
+        help=" ",
+    )
     args = argparser.parse_args()
     if args.mqtt_password_path:
         # .read_text() replaces \r\n with \n
@@ -346,10 +413,22 @@ def _main() -> None:
             mqtt_password = mqtt_password[:-1]
     else:
         mqtt_password = args.mqtt_password
+    # pylint: disable=protected-access
+    if not systemctl_mqtt._homeassistant.validate_node_id(args.homeassistant_node_id):
+        raise ValueError(
+            "invalid home assistant node id {!r} (length >= 1, allowed characters: {})".format(
+                args.homeassistant_node_id,
+                # pylint: disable=protected-access
+                systemctl_mqtt._homeassistant.NODE_ID_ALLOWED_CHARS,
+            )
+            + "\nchange --homeassistant-node-id"
+        )
     _run(
         mqtt_host=args.mqtt_host,
         mqtt_port=args.mqtt_port,
         mqtt_username=args.mqtt_username,
         mqtt_password=mqtt_password,
         mqtt_topic_prefix=args.mqtt_topic_prefix,
+        homeassistant_discovery_prefix=args.homeassistant_discovery_prefix,
+        homeassistant_node_id=args.homeassistant_node_id,
     )

systemctl_mqtt/_homeassistant.py

@@ -0,0 +1,18 @@
+import re
+
+import systemctl_mqtt._utils
+
+NODE_ID_ALLOWED_CHARS = r"a-zA-Z0-9_-"
+
+
+def get_default_node_id() -> str:
+    return re.sub(
+        r"[^{}]".format(NODE_ID_ALLOWED_CHARS),
+        "",
+        # pylint: disable=protected-access
+        systemctl_mqtt._utils.get_hostname(),
+    )
+
+
+def validate_node_id(node_id: str) -> bool:
+    return re.match(r"^[{}]+$".format(NODE_ID_ALLOWED_CHARS), node_id) is not None

systemctl_mqtt/_mqtt.py

@@ -0,0 +1,5 @@
+import json
+
+
+def encode_bool(value: bool) -> str:
+    return json.dumps(value)

systemctl_mqtt/_utils.py

@@ -0,0 +1,5 @@
+import socket
+
+
+def get_hostname() -> str:
+    return socket.gethostname()

tests/test_cli.py

@@ -21,6 +21,10 @@
 import pytest
 
 import systemctl_mqtt
+import systemctl_mqtt._homeassistant
+import systemctl_mqtt._utils
+
+# pylint: disable=protected-access
 
 
 @pytest.mark.parametrize(
@@ -100,7 +104,9 @@ def test__main(
     # pylint: disable=too-many-arguments
     with unittest.mock.patch("systemctl_mqtt._run") as run_mock, unittest.mock.patch(
         "sys.argv", argv
-    ), unittest.mock.patch("systemctl_mqtt._get_hostname", return_value="hostname"):
+    ), unittest.mock.patch(
+        "systemctl_mqtt._utils.get_hostname", return_value="hostname"
+    ):
         # pylint: disable=protected-access
         systemctl_mqtt._main()
     run_mock.assert_called_once_with(
@@ -109,6 +115,8 @@ def test__main(
         mqtt_username=expected_username,
         mqtt_password=expected_password,
         mqtt_topic_prefix=expected_topic_prefix or "systemctl/hostname",
+        homeassistant_discovery_prefix="homeassistant",
+        homeassistant_node_id="hostname",
     )
 
 
@@ -141,7 +149,9 @@ def test__main_password_file(tmpdir, password_file_content, expected_password):
             "--mqtt-password-file",
             str(mqtt_password_path),
         ],
-    ), unittest.mock.patch("systemctl_mqtt._get_hostname", return_value="hostname"):
+    ), unittest.mock.patch(
+        "systemctl_mqtt._utils.get_hostname", return_value="hostname"
+    ):
         # pylint: disable=protected-access
         systemctl_mqtt._main()
     run_mock.assert_called_once_with(
@@ -150,6 +160,8 @@ def test__main_password_file(tmpdir, password_file_content, expected_password):
         mqtt_username="me",
         mqtt_password=expected_password,
         mqtt_topic_prefix="systemctl/hostname",
+        homeassistant_discovery_prefix="homeassistant",
+        homeassistant_node_id="hostname",
     )
 
 
@@ -179,8 +191,43 @@ def test__main_password_file_collision(capsys):
     )
 
 
-@pytest.mark.parametrize("hostname", ["test"])
-def test__get_hostname(hostname):
-    with unittest.mock.patch("socket.gethostname", return_value=hostname):
-        # pylint: disable=protected-access
-        assert systemctl_mqtt._get_hostname() == hostname
+@pytest.mark.parametrize(
+    ("args", "discovery_prefix"),
+    [
+        ([], "homeassistant"),
+        (["--homeassistant-discovery-prefix", "home/assistant"], "home/assistant"),
+    ],
+)
+def test__main_homeassistant_discovery_prefix(args, discovery_prefix):
+    with unittest.mock.patch("systemctl_mqtt._run") as run_mock, unittest.mock.patch(
+        "sys.argv", ["", "--mqtt-host", "mqtt-broker.local"] + args
+    ):
+        systemctl_mqtt._main()
+    assert run_mock.call_count == 1
+    assert run_mock.call_args[1]["homeassistant_discovery_prefix"] == discovery_prefix
+
+
+@pytest.mark.parametrize(
+    ("args", "node_id"),
+    [([], "fallback"), (["--homeassistant-node-id", "raspberrypi"], "raspberrypi"),],
+)
+def test__main_homeassistant_node_id(args, node_id):
+    with unittest.mock.patch("systemctl_mqtt._run") as run_mock, unittest.mock.patch(
+        "sys.argv", ["", "--mqtt-host", "mqtt-broker.local"] + args
+    ), unittest.mock.patch(
+        "systemctl_mqtt._utils.get_hostname", return_value="fallback",
+    ):
+        systemctl_mqtt._main()
+    assert run_mock.call_count == 1
+    assert run_mock.call_args[1]["homeassistant_node_id"] == node_id
+
+
+@pytest.mark.parametrize(
+    "args", [["--homeassistant-node-id", "no pe"], ["--homeassistant-node-id", ""]],
+)
+def test__main_homeassistant_node_id_invalid(args):
+    with unittest.mock.patch(
+        "sys.argv", ["", "--mqtt-host", "mqtt-broker.local"] + args
+    ):
+        with pytest.raises(ValueError):
+            systemctl_mqtt._main()

tests/test_homeassistant.py

@@ -0,0 +1,37 @@
+import unittest.mock
+
+import pytest
+
+import systemctl_mqtt._homeassistant
+
+# pylint: disable=protected-access
+
+
+@pytest.mark.parametrize(
+    ("hostname", "expected_node_id"),
+    [
+        ("raspberrypi", "raspberrypi"),
+        ("da-sh", "da-sh"),
+        ("under_score", "under_score"),
+        ("someone evil mocked the hostname", "someoneevilmockedthehostname"),
+    ],
+)
+def test_get_default_node_id(hostname, expected_node_id):
+    with unittest.mock.patch(
+        "systemctl_mqtt._utils.get_hostname", return_value=hostname
+    ):
+        assert systemctl_mqtt._homeassistant.get_default_node_id() == expected_node_id
+
+
+@pytest.mark.parametrize(
+    ("node_id", "valid"),
+    [
+        ("raspberrypi", True),
+        ("da-sh", True),
+        ("under_score", True),
+        ('" or ""="', False),
+        ("", False),
+    ],
+)
+def test_validate_node_id(node_id, valid):
+    assert systemctl_mqtt._homeassistant.validate_node_id(node_id) == valid