1
0

__init__.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410
  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: dbus.proxies.Interface = (
  53. systemctl_mqtt._dbus.get_login_manager()
  54. )
  55. self._shutdown_lock: typing.Optional[dbus.types.UnixFd] = None
  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": f"{self._homeassistant_node_id} preparing for shutdown",
  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. class _MQTTActionLockAllSessions(_MQTTAction):
  208. def trigger(self, state: _State) -> None:
  209. # pylint: disable=protected-access
  210. systemctl_mqtt._dbus.lock_all_sessions()
  211. def __str__(self) -> str:
  212. return type(self).__name__
  213. _MQTT_TOPIC_SUFFIX_ACTION_MAPPING = {
  214. "poweroff": _MQTTActionSchedulePoweroff(),
  215. "lock-all-sessions": _MQTTActionLockAllSessions(),
  216. }
  217. def _mqtt_on_connect(
  218. mqtt_client: paho.mqtt.client.Client,
  219. state: _State,
  220. flags: typing.Dict,
  221. return_code: int,
  222. ) -> None:
  223. # pylint: disable=unused-argument; callback
  224. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L441
  225. assert return_code == 0, return_code # connection accepted
  226. mqtt_broker_host, mqtt_broker_port = mqtt_client.socket().getpeername()
  227. _LOGGER.debug("connected to MQTT broker %s:%d", mqtt_broker_host, mqtt_broker_port)
  228. if not state.shutdown_lock_acquired:
  229. state.acquire_shutdown_lock()
  230. state.register_prepare_for_shutdown_handler(mqtt_client=mqtt_client)
  231. state.publish_preparing_for_shutdown(mqtt_client=mqtt_client)
  232. state.publish_preparing_for_shutdown_homeassistant_config(mqtt_client=mqtt_client)
  233. for topic_suffix, action in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.items():
  234. topic = state.mqtt_topic_prefix + "/" + topic_suffix
  235. _LOGGER.info("subscribing to %s", topic)
  236. mqtt_client.subscribe(topic)
  237. mqtt_client.message_callback_add(
  238. sub=topic, callback=action.mqtt_message_callback
  239. )
  240. _LOGGER.debug(
  241. "registered MQTT callback for topic %s triggering %s", topic, action
  242. )
  243. def _run(
  244. *,
  245. mqtt_host: str,
  246. mqtt_port: int,
  247. mqtt_username: typing.Optional[str],
  248. mqtt_password: typing.Optional[str],
  249. mqtt_topic_prefix: str,
  250. homeassistant_discovery_prefix: str,
  251. homeassistant_node_id: str,
  252. poweroff_delay: datetime.timedelta,
  253. mqtt_disable_tls: bool = False,
  254. ) -> None:
  255. # pylint: disable=too-many-arguments
  256. # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
  257. dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
  258. # https://pypi.org/project/paho-mqtt/
  259. mqtt_client = paho.mqtt.client.Client(
  260. userdata=_State(
  261. mqtt_topic_prefix=mqtt_topic_prefix,
  262. homeassistant_discovery_prefix=homeassistant_discovery_prefix,
  263. homeassistant_node_id=homeassistant_node_id,
  264. poweroff_delay=poweroff_delay,
  265. )
  266. )
  267. mqtt_client.on_connect = _mqtt_on_connect
  268. if not mqtt_disable_tls:
  269. mqtt_client.tls_set(ca_certs=None) # enable tls trusting default system certs
  270. _LOGGER.info(
  271. "connecting to MQTT broker %s:%d (TLS %s)",
  272. mqtt_host,
  273. mqtt_port,
  274. "disabled" if mqtt_disable_tls else "enabled",
  275. )
  276. if mqtt_username:
  277. mqtt_client.username_pw_set(username=mqtt_username, password=mqtt_password)
  278. elif mqtt_password:
  279. raise ValueError("Missing MQTT username")
  280. mqtt_client.connect(host=mqtt_host, port=mqtt_port)
  281. # loop_start runs loop_forever in a new thread (daemon)
  282. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1814
  283. # loop_forever attempts to reconnect if disconnected
  284. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1744
  285. mqtt_client.loop_start()
  286. try:
  287. gi.repository.GLib.MainLoop().run()
  288. finally:
  289. # blocks until loop_forever stops
  290. _LOGGER.debug("waiting for MQTT loop to stop")
  291. mqtt_client.loop_stop()
  292. _LOGGER.debug("MQTT loop stopped")
  293. def _main() -> None:
  294. logging.basicConfig(
  295. level=logging.DEBUG,
  296. format="%(asctime)s:%(levelname)s:%(message)s",
  297. datefmt="%Y-%m-%dT%H:%M:%S%z",
  298. )
  299. argparser = argparse.ArgumentParser(
  300. description="MQTT client triggering & reporting shutdown on systemd-based systems",
  301. formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  302. )
  303. argparser.add_argument("--mqtt-host", type=str, required=True)
  304. argparser.add_argument(
  305. "--mqtt-port",
  306. type=int,
  307. help=f"default {_MQTT_DEFAULT_TLS_PORT} ({_MQTT_DEFAULT_PORT} with --mqtt-disable-tls)",
  308. )
  309. argparser.add_argument("--mqtt-username", type=str)
  310. argparser.add_argument("--mqtt-disable-tls", action="store_true")
  311. password_argument_group = argparser.add_mutually_exclusive_group()
  312. password_argument_group.add_argument("--mqtt-password", type=str)
  313. password_argument_group.add_argument(
  314. "--mqtt-password-file",
  315. type=pathlib.Path,
  316. metavar="PATH",
  317. dest="mqtt_password_path",
  318. help="stripping trailing newline",
  319. )
  320. argparser.add_argument(
  321. "--mqtt-topic-prefix",
  322. type=str,
  323. # pylint: disable=protected-access
  324. default="systemctl/" + systemctl_mqtt._utils.get_hostname(),
  325. help=" ", # show default
  326. )
  327. # https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix
  328. argparser.add_argument(
  329. "--homeassistant-discovery-prefix", type=str, default="homeassistant", help=" "
  330. )
  331. argparser.add_argument(
  332. "--homeassistant-node-id",
  333. type=str,
  334. # pylint: disable=protected-access
  335. default=systemctl_mqtt._homeassistant.get_default_node_id(),
  336. help=" ",
  337. )
  338. argparser.add_argument(
  339. "--poweroff-delay-seconds", type=float, default=4.0, help=" "
  340. )
  341. args = argparser.parse_args()
  342. if args.mqtt_port:
  343. mqtt_port = args.mqtt_port
  344. elif args.mqtt_disable_tls:
  345. mqtt_port = _MQTT_DEFAULT_PORT
  346. else:
  347. mqtt_port = _MQTT_DEFAULT_TLS_PORT
  348. if args.mqtt_password_path:
  349. # .read_text() replaces \r\n with \n
  350. mqtt_password = args.mqtt_password_path.read_bytes().decode()
  351. if mqtt_password.endswith("\r\n"):
  352. mqtt_password = mqtt_password[:-2]
  353. elif mqtt_password.endswith("\n"):
  354. mqtt_password = mqtt_password[:-1]
  355. else:
  356. mqtt_password = args.mqtt_password
  357. # pylint: disable=protected-access
  358. if not systemctl_mqtt._homeassistant.validate_node_id(args.homeassistant_node_id):
  359. raise ValueError(
  360. # pylint: disable=protected-access
  361. f"invalid home assistant node id {args.homeassistant_node_id!r} (length >= 1"
  362. f", allowed characters: {systemctl_mqtt._homeassistant.NODE_ID_ALLOWED_CHARS})"
  363. "\nchange --homeassistant-node-id"
  364. )
  365. _run(
  366. mqtt_host=args.mqtt_host,
  367. mqtt_port=mqtt_port,
  368. mqtt_disable_tls=args.mqtt_disable_tls,
  369. mqtt_username=args.mqtt_username,
  370. mqtt_password=mqtt_password,
  371. mqtt_topic_prefix=args.mqtt_topic_prefix,
  372. homeassistant_discovery_prefix=args.homeassistant_discovery_prefix,
  373. homeassistant_node_id=args.homeassistant_node_id,
  374. poweroff_delay=datetime.timedelta(seconds=args.poweroff_delay_seconds),
  375. )