__init__.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402
  1. # systemctl-mqtt - MQTT client triggering & reporting shutdown on systemd-based systems
  2. #
  3. # Copyright (C) 2020 Fabian Peter Hammerle <fabian@hammerle.me>
  4. #
  5. # This program is free software: you can redistribute it and/or modify
  6. # it under the terms of the GNU General Public License as published by
  7. # the Free Software Foundation, either version 3 of the License, or
  8. # any later version.
  9. #
  10. # This program is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. # GNU General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU General Public License
  16. # along with this program. If not, see <https://www.gnu.org/licenses/>.
  17. import abc
  18. import argparse
  19. import datetime
  20. import functools
  21. import json
  22. import logging
  23. import os
  24. import pathlib
  25. import socket
  26. import threading
  27. import typing
  28. import dbus
  29. import dbus.mainloop.glib
  30. # black keeps inserting a blank line above
  31. # https://pygobject.readthedocs.io/en/latest/getting_started.html#ubuntu-logo-ubuntu-debian-logo-debian
  32. import gi.repository.GLib # pylint-import-requirements: imports=PyGObject
  33. import paho.mqtt.client
  34. import systemctl_mqtt._dbus
  35. import systemctl_mqtt._homeassistant
  36. import systemctl_mqtt._mqtt
  37. _MQTT_DEFAULT_PORT = 1883
  38. _MQTT_DEFAULT_TLS_PORT = 8883
  39. _LOGGER = logging.getLogger(__name__)
  40. class _State:
  41. def __init__(
  42. self,
  43. *,
  44. mqtt_topic_prefix: str,
  45. homeassistant_discovery_prefix: str,
  46. homeassistant_node_id: str,
  47. poweroff_delay: datetime.timedelta,
  48. ) -> None:
  49. self._mqtt_topic_prefix = mqtt_topic_prefix
  50. self._homeassistant_discovery_prefix = homeassistant_discovery_prefix
  51. self._homeassistant_node_id = homeassistant_node_id
  52. self._login_manager = (
  53. systemctl_mqtt._dbus.get_login_manager()
  54. ) # type: dbus.proxies.Interface
  55. self._shutdown_lock = None # type: typing.Optional[dbus.types.UnixFd]
  56. self._shutdown_lock_mutex = threading.Lock()
  57. self.poweroff_delay = poweroff_delay
  58. @property
  59. def mqtt_topic_prefix(self) -> str:
  60. return self._mqtt_topic_prefix
  61. @property
  62. def shutdown_lock_acquired(self) -> bool:
  63. return self._shutdown_lock is not None
  64. def acquire_shutdown_lock(self) -> None:
  65. with self._shutdown_lock_mutex:
  66. assert self._shutdown_lock is None
  67. # https://www.freedesktop.org/wiki/Software/systemd/inhibit/
  68. self._shutdown_lock = self._login_manager.Inhibit(
  69. "shutdown", "systemctl-mqtt", "Report shutdown via MQTT", "delay",
  70. )
  71. _LOGGER.debug("acquired shutdown inhibitor lock")
  72. def release_shutdown_lock(self) -> None:
  73. with self._shutdown_lock_mutex:
  74. if self._shutdown_lock:
  75. # https://dbus.freedesktop.org/doc/dbus-python/dbus.types.html#dbus.types.UnixFd.take
  76. os.close(self._shutdown_lock.take())
  77. _LOGGER.debug("released shutdown inhibitor lock")
  78. self._shutdown_lock = None
  79. @property
  80. def _preparing_for_shutdown_topic(self) -> str:
  81. return self.mqtt_topic_prefix + "/preparing-for-shutdown"
  82. def _publish_preparing_for_shutdown(
  83. self, mqtt_client: paho.mqtt.client.Client, active: bool, block: bool,
  84. ) -> None:
  85. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1199
  86. topic = self._preparing_for_shutdown_topic
  87. # pylint: disable=protected-access
  88. payload = systemctl_mqtt._mqtt.encode_bool(active)
  89. _LOGGER.info("publishing %r on %s", payload, topic)
  90. msg_info = mqtt_client.publish(
  91. topic=topic, payload=payload, retain=True,
  92. ) # type: paho.mqtt.client.MQTTMessageInfo
  93. if not block:
  94. return
  95. msg_info.wait_for_publish()
  96. if msg_info.rc != paho.mqtt.client.MQTT_ERR_SUCCESS:
  97. _LOGGER.error(
  98. "failed to publish on %s (return code %d)", topic, msg_info.rc
  99. )
  100. def _prepare_for_shutdown_handler(
  101. self, active: dbus.Boolean, mqtt_client: paho.mqtt.client.Client
  102. ) -> None:
  103. assert isinstance(active, dbus.Boolean)
  104. active = bool(active)
  105. self._publish_preparing_for_shutdown(
  106. mqtt_client=mqtt_client, active=active, block=True,
  107. )
  108. if active:
  109. self.release_shutdown_lock()
  110. else:
  111. self.acquire_shutdown_lock()
  112. def register_prepare_for_shutdown_handler(
  113. self, mqtt_client: paho.mqtt.client.Client
  114. ) -> None:
  115. self._login_manager.connect_to_signal(
  116. signal_name="PrepareForShutdown",
  117. handler_function=functools.partial(
  118. self._prepare_for_shutdown_handler, mqtt_client=mqtt_client
  119. ),
  120. )
  121. def publish_preparing_for_shutdown(
  122. self, mqtt_client: paho.mqtt.client.Client,
  123. ) -> None:
  124. try:
  125. active = self._login_manager.Get(
  126. "org.freedesktop.login1.Manager",
  127. "PreparingForShutdown",
  128. dbus_interface="org.freedesktop.DBus.Properties",
  129. )
  130. except dbus.DBusException as exc:
  131. _LOGGER.error(
  132. "failed to read logind's PreparingForShutdown property: %s",
  133. exc.get_dbus_message(),
  134. )
  135. return
  136. assert isinstance(active, dbus.Boolean), active
  137. self._publish_preparing_for_shutdown(
  138. mqtt_client=mqtt_client,
  139. active=bool(active),
  140. # https://github.com/eclipse/paho.mqtt.python/issues/439#issuecomment-565514393
  141. block=False,
  142. )
  143. def publish_preparing_for_shutdown_homeassistant_config(
  144. self, mqtt_client: paho.mqtt.client.Client
  145. ) -> None:
  146. # <discovery_prefix>/<component>/[<node_id>/]<object_id>/config
  147. # https://www.home-assistant.io/docs/mqtt/discovery/
  148. discovery_topic = "/".join(
  149. (
  150. self._homeassistant_discovery_prefix,
  151. "binary_sensor",
  152. self._homeassistant_node_id,
  153. "preparing-for-shutdown",
  154. "config",
  155. )
  156. )
  157. unique_id = "/".join(
  158. (
  159. "systemctl-mqtt",
  160. self._homeassistant_node_id,
  161. "logind",
  162. "preparing-for-shutdown",
  163. )
  164. )
  165. # https://www.home-assistant.io/integrations/binary_sensor.mqtt/#configuration-variables
  166. config = {
  167. "unique_id": unique_id,
  168. "state_topic": self._preparing_for_shutdown_topic,
  169. # pylint: disable=protected-access
  170. "payload_on": systemctl_mqtt._mqtt.encode_bool(True),
  171. "payload_off": systemctl_mqtt._mqtt.encode_bool(False),
  172. # friendly_name & template for default entity_id
  173. "name": "{} preparing for shutdown".format(self._homeassistant_node_id),
  174. }
  175. _LOGGER.debug("publishing home assistant config on %s", discovery_topic)
  176. mqtt_client.publish(
  177. topic=discovery_topic, payload=json.dumps(config), retain=True,
  178. )
  179. class _MQTTAction(metaclass=abc.ABCMeta):
  180. @abc.abstractmethod
  181. def trigger(self, state: _State) -> None:
  182. pass # pragma: no cover
  183. def mqtt_message_callback(
  184. self,
  185. mqtt_client: paho.mqtt.client.Client,
  186. state: _State,
  187. message: paho.mqtt.client.MQTTMessage,
  188. ) -> None:
  189. # pylint: disable=unused-argument; callback
  190. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L3416
  191. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L469
  192. _LOGGER.debug("received topic=%s payload=%r", message.topic, message.payload)
  193. if message.retain:
  194. _LOGGER.info("ignoring retained message")
  195. return
  196. _LOGGER.debug("executing action %s", self)
  197. self.trigger(state=state)
  198. _LOGGER.debug("completed action %s", self)
  199. class _MQTTActionSchedulePoweroff(_MQTTAction):
  200. def trigger(self, state: _State) -> None:
  201. # pylint: disable=protected-access
  202. systemctl_mqtt._dbus.schedule_shutdown(
  203. action="poweroff", delay=state.poweroff_delay
  204. )
  205. def __str__(self) -> str:
  206. return type(self).__name__
  207. _MQTT_TOPIC_SUFFIX_ACTION_MAPPING = {"poweroff": _MQTTActionSchedulePoweroff()}
  208. def _mqtt_on_connect(
  209. mqtt_client: paho.mqtt.client.Client,
  210. state: _State,
  211. flags: typing.Dict,
  212. return_code: int,
  213. ) -> None:
  214. # pylint: disable=unused-argument; callback
  215. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L441
  216. assert return_code == 0, return_code # connection accepted
  217. mqtt_broker_host, mqtt_broker_port = mqtt_client.socket().getpeername()
  218. _LOGGER.debug("connected to MQTT broker %s:%d", mqtt_broker_host, mqtt_broker_port)
  219. if not state.shutdown_lock_acquired:
  220. state.acquire_shutdown_lock()
  221. state.register_prepare_for_shutdown_handler(mqtt_client=mqtt_client)
  222. state.publish_preparing_for_shutdown(mqtt_client=mqtt_client)
  223. state.publish_preparing_for_shutdown_homeassistant_config(mqtt_client=mqtt_client)
  224. for topic_suffix, action in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.items():
  225. topic = state.mqtt_topic_prefix + "/" + topic_suffix
  226. _LOGGER.info("subscribing to %s", topic)
  227. mqtt_client.subscribe(topic)
  228. mqtt_client.message_callback_add(
  229. sub=topic, callback=action.mqtt_message_callback
  230. )
  231. _LOGGER.debug(
  232. "registered MQTT callback for topic %s triggering %s", topic, action
  233. )
  234. def _run(
  235. *,
  236. mqtt_host: str,
  237. mqtt_port: int,
  238. mqtt_username: typing.Optional[str],
  239. mqtt_password: typing.Optional[str],
  240. mqtt_topic_prefix: str,
  241. homeassistant_discovery_prefix: str,
  242. homeassistant_node_id: str,
  243. poweroff_delay: datetime.timedelta,
  244. mqtt_disable_tls: bool = False,
  245. ) -> None:
  246. # pylint: disable=too-many-arguments
  247. # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
  248. dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
  249. # https://pypi.org/project/paho-mqtt/
  250. mqtt_client = paho.mqtt.client.Client(
  251. userdata=_State(
  252. mqtt_topic_prefix=mqtt_topic_prefix,
  253. homeassistant_discovery_prefix=homeassistant_discovery_prefix,
  254. homeassistant_node_id=homeassistant_node_id,
  255. poweroff_delay=poweroff_delay,
  256. )
  257. )
  258. mqtt_client.on_connect = _mqtt_on_connect
  259. if not mqtt_disable_tls:
  260. mqtt_client.tls_set(ca_certs=None) # enable tls trusting default system certs
  261. _LOGGER.info(
  262. "connecting to MQTT broker %s:%d (TLS %s)",
  263. mqtt_host,
  264. mqtt_port,
  265. "disabled" if mqtt_disable_tls else "enabled",
  266. )
  267. if mqtt_username:
  268. mqtt_client.username_pw_set(username=mqtt_username, password=mqtt_password)
  269. elif mqtt_password:
  270. raise ValueError("Missing MQTT username")
  271. mqtt_client.connect(host=mqtt_host, port=mqtt_port)
  272. # loop_start runs loop_forever in a new thread (daemon)
  273. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1814
  274. # loop_forever attempts to reconnect if disconnected
  275. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1744
  276. mqtt_client.loop_start()
  277. try:
  278. gi.repository.GLib.MainLoop().run()
  279. finally:
  280. # blocks until loop_forever stops
  281. _LOGGER.debug("waiting for MQTT loop to stop")
  282. mqtt_client.loop_stop()
  283. _LOGGER.debug("MQTT loop stopped")
  284. def _main() -> None:
  285. logging.basicConfig(
  286. level=logging.DEBUG,
  287. format="%(asctime)s:%(levelname)s:%(message)s",
  288. datefmt="%Y-%m-%dT%H:%M:%S%z",
  289. )
  290. argparser = argparse.ArgumentParser(
  291. description="MQTT client triggering & reporting shutdown on systemd-based systems",
  292. formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  293. )
  294. argparser.add_argument("--mqtt-host", type=str, required=True)
  295. argparser.add_argument(
  296. "--mqtt-port",
  297. type=int,
  298. help="default {} ({} with --mqtt-disable-tls)".format(
  299. _MQTT_DEFAULT_TLS_PORT, _MQTT_DEFAULT_PORT
  300. ),
  301. )
  302. argparser.add_argument("--mqtt-username", type=str)
  303. argparser.add_argument("--mqtt-disable-tls", action="store_true")
  304. password_argument_group = argparser.add_mutually_exclusive_group()
  305. password_argument_group.add_argument("--mqtt-password", type=str)
  306. password_argument_group.add_argument(
  307. "--mqtt-password-file",
  308. type=pathlib.Path,
  309. metavar="PATH",
  310. dest="mqtt_password_path",
  311. help="stripping trailing newline",
  312. )
  313. argparser.add_argument(
  314. "--mqtt-topic-prefix",
  315. type=str,
  316. # pylint: disable=protected-access
  317. default="systemctl/" + systemctl_mqtt._utils.get_hostname(),
  318. help=" ", # show default
  319. )
  320. # https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix
  321. argparser.add_argument(
  322. "--homeassistant-discovery-prefix", type=str, default="homeassistant", help=" ",
  323. )
  324. argparser.add_argument(
  325. "--homeassistant-node-id",
  326. type=str,
  327. # pylint: disable=protected-access
  328. default=systemctl_mqtt._homeassistant.get_default_node_id(),
  329. help=" ",
  330. )
  331. argparser.add_argument(
  332. "--poweroff-delay-seconds", type=float, default=4.0, help=" ",
  333. )
  334. args = argparser.parse_args()
  335. if args.mqtt_port:
  336. mqtt_port = args.mqtt_port
  337. elif args.mqtt_disable_tls:
  338. mqtt_port = _MQTT_DEFAULT_PORT
  339. else:
  340. mqtt_port = _MQTT_DEFAULT_TLS_PORT
  341. if args.mqtt_password_path:
  342. # .read_text() replaces \r\n with \n
  343. mqtt_password = args.mqtt_password_path.read_bytes().decode()
  344. if mqtt_password.endswith("\r\n"):
  345. mqtt_password = mqtt_password[:-2]
  346. elif mqtt_password.endswith("\n"):
  347. mqtt_password = mqtt_password[:-1]
  348. else:
  349. mqtt_password = args.mqtt_password
  350. # pylint: disable=protected-access
  351. if not systemctl_mqtt._homeassistant.validate_node_id(args.homeassistant_node_id):
  352. raise ValueError(
  353. "invalid home assistant node id {!r} (length >= 1, allowed characters: {})".format(
  354. args.homeassistant_node_id,
  355. # pylint: disable=protected-access
  356. systemctl_mqtt._homeassistant.NODE_ID_ALLOWED_CHARS,
  357. )
  358. + "\nchange --homeassistant-node-id"
  359. )
  360. _run(
  361. mqtt_host=args.mqtt_host,
  362. mqtt_port=mqtt_port,
  363. mqtt_disable_tls=args.mqtt_disable_tls,
  364. mqtt_username=args.mqtt_username,
  365. mqtt_password=mqtt_password,
  366. mqtt_topic_prefix=args.mqtt_topic_prefix,
  367. homeassistant_discovery_prefix=args.homeassistant_discovery_prefix,
  368. homeassistant_node_id=args.homeassistant_node_id,
  369. poweroff_delay=datetime.timedelta(seconds=args.poweroff_delay_seconds),
  370. )