1
0

__init__.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412
  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. mqtt_topic_prefix: str,
  44. homeassistant_discovery_prefix: str,
  45. homeassistant_node_id: str,
  46. poweroff_delay: datetime.timedelta,
  47. ) -> None:
  48. self._mqtt_topic_prefix = mqtt_topic_prefix
  49. self._homeassistant_discovery_prefix = homeassistant_discovery_prefix
  50. self._homeassistant_node_id = homeassistant_node_id
  51. self._login_manager = (
  52. systemctl_mqtt._dbus.get_login_manager()
  53. ) # type: dbus.proxies.Interface
  54. self._shutdown_lock = None # type: typing.Optional[dbus.types.UnixFd]
  55. self._shutdown_lock_mutex = threading.Lock()
  56. self.poweroff_delay = poweroff_delay
  57. @property
  58. def mqtt_topic_prefix(self) -> str:
  59. return self._mqtt_topic_prefix
  60. @property
  61. def shutdown_lock_acquired(self) -> bool:
  62. return self._shutdown_lock is not None
  63. def acquire_shutdown_lock(self) -> None:
  64. with self._shutdown_lock_mutex:
  65. assert self._shutdown_lock is None
  66. # https://www.freedesktop.org/wiki/Software/systemd/inhibit/
  67. self._shutdown_lock = self._login_manager.Inhibit(
  68. "shutdown", "systemctl-mqtt", "Report shutdown via MQTT", "delay"
  69. )
  70. _LOGGER.debug("acquired shutdown inhibitor lock")
  71. def release_shutdown_lock(self) -> None:
  72. with self._shutdown_lock_mutex:
  73. if self._shutdown_lock:
  74. # https://dbus.freedesktop.org/doc/dbus-python/dbus.types.html#dbus.types.UnixFd.take
  75. os.close(self._shutdown_lock.take())
  76. _LOGGER.debug("released shutdown inhibitor lock")
  77. self._shutdown_lock = None
  78. @property
  79. def _preparing_for_shutdown_topic(self) -> str:
  80. return self.mqtt_topic_prefix + "/preparing-for-shutdown"
  81. def _publish_preparing_for_shutdown(
  82. self, mqtt_client: paho.mqtt.client.Client, active: bool, block: bool
  83. ) -> None:
  84. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1199
  85. topic = self._preparing_for_shutdown_topic
  86. # pylint: disable=protected-access
  87. payload = systemctl_mqtt._mqtt.encode_bool(active)
  88. _LOGGER.info("publishing %r on %s", payload, topic)
  89. msg_info = mqtt_client.publish(
  90. topic=topic, payload=payload, retain=True
  91. ) # type: paho.mqtt.client.MQTTMessageInfo
  92. if not block:
  93. return
  94. msg_info.wait_for_publish()
  95. if msg_info.rc != paho.mqtt.client.MQTT_ERR_SUCCESS:
  96. _LOGGER.error(
  97. "failed to publish on %s (return code %d)", topic, msg_info.rc
  98. )
  99. def _prepare_for_shutdown_handler(
  100. self, active: dbus.Boolean, mqtt_client: paho.mqtt.client.Client
  101. ) -> None:
  102. assert isinstance(active, dbus.Boolean)
  103. active = bool(active)
  104. self._publish_preparing_for_shutdown(
  105. mqtt_client=mqtt_client, active=active, block=True
  106. )
  107. if active:
  108. self.release_shutdown_lock()
  109. else:
  110. self.acquire_shutdown_lock()
  111. def register_prepare_for_shutdown_handler(
  112. self, mqtt_client: paho.mqtt.client.Client
  113. ) -> None:
  114. self._login_manager.connect_to_signal(
  115. signal_name="PrepareForShutdown",
  116. handler_function=functools.partial(
  117. self._prepare_for_shutdown_handler, mqtt_client=mqtt_client
  118. ),
  119. )
  120. def publish_preparing_for_shutdown(
  121. self, mqtt_client: paho.mqtt.client.Client
  122. ) -> None:
  123. try:
  124. active = self._login_manager.Get(
  125. "org.freedesktop.login1.Manager",
  126. "PreparingForShutdown",
  127. dbus_interface="org.freedesktop.DBus.Properties",
  128. )
  129. except dbus.DBusException as exc:
  130. _LOGGER.error(
  131. "failed to read logind's PreparingForShutdown property: %s",
  132. exc.get_dbus_message(),
  133. )
  134. return
  135. assert isinstance(active, dbus.Boolean), active
  136. self._publish_preparing_for_shutdown(
  137. mqtt_client=mqtt_client,
  138. active=bool(active),
  139. # https://github.com/eclipse/paho.mqtt.python/issues/439#issuecomment-565514393
  140. block=False,
  141. )
  142. def publish_preparing_for_shutdown_homeassistant_config(
  143. self, mqtt_client: paho.mqtt.client.Client
  144. ) -> None:
  145. # <discovery_prefix>/<component>/[<node_id>/]<object_id>/config
  146. # https://www.home-assistant.io/docs/mqtt/discovery/
  147. discovery_topic = "/".join(
  148. (
  149. self._homeassistant_discovery_prefix,
  150. "binary_sensor",
  151. self._homeassistant_node_id,
  152. "preparing-for-shutdown",
  153. "config",
  154. )
  155. )
  156. unique_id = "/".join(
  157. (
  158. "systemctl-mqtt",
  159. self._homeassistant_node_id,
  160. "logind",
  161. "preparing-for-shutdown",
  162. )
  163. )
  164. # https://www.home-assistant.io/integrations/binary_sensor.mqtt/#configuration-variables
  165. config = {
  166. "unique_id": unique_id,
  167. "state_topic": self._preparing_for_shutdown_topic,
  168. # pylint: disable=protected-access
  169. "payload_on": systemctl_mqtt._mqtt.encode_bool(True),
  170. "payload_off": systemctl_mqtt._mqtt.encode_bool(False),
  171. # friendly_name & template for default entity_id
  172. "name": "{} preparing for shutdown".format(self._homeassistant_node_id),
  173. }
  174. _LOGGER.debug("publishing home assistant config on %s", discovery_topic)
  175. mqtt_client.publish(
  176. topic=discovery_topic, payload=json.dumps(config), retain=True
  177. )
  178. class _MQTTAction(metaclass=abc.ABCMeta):
  179. @abc.abstractmethod
  180. def trigger(self, state: _State) -> None:
  181. pass # pragma: no cover
  182. def mqtt_message_callback(
  183. self,
  184. mqtt_client: paho.mqtt.client.Client,
  185. state: _State,
  186. message: paho.mqtt.client.MQTTMessage,
  187. ) -> None:
  188. # pylint: disable=unused-argument; callback
  189. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L3416
  190. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L469
  191. _LOGGER.debug("received topic=%s payload=%r", message.topic, message.payload)
  192. if message.retain:
  193. _LOGGER.info("ignoring retained message")
  194. return
  195. _LOGGER.debug("executing action %s", self)
  196. self.trigger(state=state)
  197. _LOGGER.debug("completed action %s", self)
  198. class _MQTTActionSchedulePoweroff(_MQTTAction):
  199. def trigger(self, state: _State) -> None:
  200. # pylint: disable=protected-access
  201. systemctl_mqtt._dbus.schedule_shutdown(
  202. action="poweroff", delay=state.poweroff_delay
  203. )
  204. def __str__(self) -> str:
  205. return type(self).__name__
  206. class _MQTTActionLockAllSessions(_MQTTAction):
  207. def trigger(self, state: _State) -> None:
  208. # pylint: disable=protected-access
  209. systemctl_mqtt._dbus.lock_all_sessions()
  210. def __str__(self) -> str:
  211. return type(self).__name__
  212. _MQTT_TOPIC_SUFFIX_ACTION_MAPPING = {
  213. "poweroff": _MQTTActionSchedulePoweroff(),
  214. "lock-all-sessions": _MQTTActionLockAllSessions(),
  215. }
  216. def _mqtt_on_connect(
  217. mqtt_client: paho.mqtt.client.Client,
  218. state: _State,
  219. flags: typing.Dict,
  220. return_code: int,
  221. ) -> None:
  222. # pylint: disable=unused-argument; callback
  223. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L441
  224. assert return_code == 0, return_code # connection accepted
  225. mqtt_broker_host, mqtt_broker_port = mqtt_client.socket().getpeername()
  226. _LOGGER.debug("connected to MQTT broker %s:%d", mqtt_broker_host, mqtt_broker_port)
  227. if not state.shutdown_lock_acquired:
  228. state.acquire_shutdown_lock()
  229. state.register_prepare_for_shutdown_handler(mqtt_client=mqtt_client)
  230. state.publish_preparing_for_shutdown(mqtt_client=mqtt_client)
  231. state.publish_preparing_for_shutdown_homeassistant_config(mqtt_client=mqtt_client)
  232. for topic_suffix, action in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.items():
  233. topic = state.mqtt_topic_prefix + "/" + topic_suffix
  234. _LOGGER.info("subscribing to %s", topic)
  235. mqtt_client.subscribe(topic)
  236. mqtt_client.message_callback_add(
  237. sub=topic, callback=action.mqtt_message_callback
  238. )
  239. _LOGGER.debug(
  240. "registered MQTT callback for topic %s triggering %s", topic, action
  241. )
  242. def _run(
  243. mqtt_host: str,
  244. mqtt_port: int,
  245. mqtt_username: typing.Optional[str],
  246. mqtt_password: typing.Optional[str],
  247. mqtt_topic_prefix: str,
  248. homeassistant_discovery_prefix: str,
  249. homeassistant_node_id: str,
  250. poweroff_delay: datetime.timedelta,
  251. mqtt_disable_tls: bool = False,
  252. ) -> None:
  253. # pylint: disable=too-many-arguments
  254. # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
  255. dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
  256. # https://pypi.org/project/paho-mqtt/
  257. mqtt_client = paho.mqtt.client.Client(
  258. userdata=_State(
  259. mqtt_topic_prefix=mqtt_topic_prefix,
  260. homeassistant_discovery_prefix=homeassistant_discovery_prefix,
  261. homeassistant_node_id=homeassistant_node_id,
  262. poweroff_delay=poweroff_delay,
  263. )
  264. )
  265. mqtt_client.on_connect = _mqtt_on_connect
  266. if not mqtt_disable_tls:
  267. mqtt_client.tls_set(ca_certs=None) # enable tls trusting default system certs
  268. _LOGGER.info(
  269. "connecting to MQTT broker %s:%d (TLS %s)",
  270. mqtt_host,
  271. mqtt_port,
  272. "disabled" if mqtt_disable_tls else "enabled",
  273. )
  274. if mqtt_username:
  275. mqtt_client.username_pw_set(username=mqtt_username, password=mqtt_password)
  276. elif mqtt_password:
  277. raise ValueError("Missing MQTT username")
  278. mqtt_client.connect(host=mqtt_host, port=mqtt_port)
  279. # loop_start runs loop_forever in a new thread (daemon)
  280. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1814
  281. # loop_forever attempts to reconnect if disconnected
  282. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1744
  283. mqtt_client.loop_start()
  284. try:
  285. gi.repository.GLib.MainLoop().run()
  286. finally:
  287. # blocks until loop_forever stops
  288. _LOGGER.debug("waiting for MQTT loop to stop")
  289. mqtt_client.loop_stop()
  290. _LOGGER.debug("MQTT loop stopped")
  291. def _main() -> None:
  292. logging.basicConfig(
  293. level=logging.DEBUG,
  294. format="%(asctime)s:%(levelname)s:%(message)s",
  295. datefmt="%Y-%m-%dT%H:%M:%S%z",
  296. )
  297. argparser = argparse.ArgumentParser(
  298. description="MQTT client triggering & reporting shutdown on systemd-based systems",
  299. formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  300. )
  301. argparser.add_argument("--mqtt-host", type=str, required=True)
  302. argparser.add_argument(
  303. "--mqtt-port",
  304. type=int,
  305. help="default {} ({} with --mqtt-disable-tls)".format(
  306. _MQTT_DEFAULT_TLS_PORT, _MQTT_DEFAULT_PORT
  307. ),
  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. "invalid home assistant node id {!r} (length >= 1, allowed characters: {})".format(
  361. args.homeassistant_node_id,
  362. # pylint: disable=protected-access
  363. systemctl_mqtt._homeassistant.NODE_ID_ALLOWED_CHARS,
  364. )
  365. + "\nchange --homeassistant-node-id"
  366. )
  367. _run(
  368. mqtt_host=args.mqtt_host,
  369. mqtt_port=mqtt_port,
  370. mqtt_disable_tls=args.mqtt_disable_tls,
  371. mqtt_username=args.mqtt_username,
  372. mqtt_password=mqtt_password,
  373. mqtt_topic_prefix=args.mqtt_topic_prefix,
  374. homeassistant_discovery_prefix=args.homeassistant_discovery_prefix,
  375. homeassistant_node_id=args.homeassistant_node_id,
  376. poweroff_delay=datetime.timedelta(seconds=args.poweroff_delay_seconds),
  377. )