__init__.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681
  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 asyncio
  20. import datetime
  21. import functools
  22. import importlib.metadata
  23. import json
  24. import logging
  25. import os
  26. import pathlib
  27. import socket
  28. import ssl
  29. import threading
  30. import typing
  31. import aiomqtt
  32. import jeepney
  33. import jeepney.bus_messages
  34. import jeepney.io.asyncio
  35. import systemctl_mqtt._dbus.login_manager
  36. import systemctl_mqtt._dbus.service_manager
  37. import systemctl_mqtt._homeassistant
  38. import systemctl_mqtt._mqtt
  39. _MQTT_DEFAULT_PORT = 1883
  40. _MQTT_DEFAULT_TLS_PORT = 8883
  41. # > payload_not_available string (Optional, default: offline)
  42. # https://web.archive.org/web/20250101075341/https://www.home-assistant.io/integrations/sensor.mqtt/#payload_not_available
  43. _MQTT_PAYLOAD_NOT_AVAILABLE = "offline"
  44. _MQTT_PAYLOAD_AVAILABLE = "online"
  45. _ARGUMENT_LOG_LEVEL_MAPPING = {
  46. a: getattr(logging, a.upper())
  47. for a in ("debug", "info", "warning", "error", "critical")
  48. }
  49. _LOGGER = logging.getLogger(__name__)
  50. class _State:
  51. # pylint: disable=too-many-instance-attributes
  52. def __init__( # pylint: disable=too-many-arguments
  53. self,
  54. *,
  55. mqtt_topic_prefix: str,
  56. homeassistant_discovery_prefix: str,
  57. homeassistant_discovery_object_id: str,
  58. poweroff_delay: datetime.timedelta,
  59. monitored_system_unit_names: typing.List[str],
  60. controlled_system_unit_names: typing.List[str],
  61. ) -> None:
  62. self._mqtt_topic_prefix = mqtt_topic_prefix
  63. self._homeassistant_discovery_prefix = homeassistant_discovery_prefix
  64. self._homeassistant_discovery_object_id = homeassistant_discovery_object_id
  65. self._login_manager = (
  66. systemctl_mqtt._dbus.login_manager.get_login_manager_proxy()
  67. )
  68. self._shutdown_lock: typing.Optional[jeepney.fds.FileDescriptor] = None
  69. self._shutdown_lock_mutex = threading.Lock()
  70. self.poweroff_delay = poweroff_delay
  71. self._monitored_system_unit_names = monitored_system_unit_names
  72. self._controlled_system_unit_names = controlled_system_unit_names
  73. @property
  74. def mqtt_topic_prefix(self) -> str:
  75. return self._mqtt_topic_prefix
  76. @property
  77. def mqtt_availability_topic(self) -> str:
  78. # > mqtt.ATTR_TOPIC: "homeassistant/status",
  79. # https://github.com/home-assistant/core/blob/2024.12.5/tests/components/mqtt/conftest.py#L23
  80. # > _MQTT_AVAILABILITY_TOPIC = "switchbot-mqtt/status"
  81. # https://github.com/fphammerle/switchbot-mqtt/blob/v3.3.1/switchbot_mqtt/__init__.py#L30
  82. return self._mqtt_topic_prefix + "/status"
  83. def get_system_unit_active_state_mqtt_topic(self, *, unit_name: str) -> str:
  84. return self._mqtt_topic_prefix + "/unit/system/" + unit_name + "/active-state"
  85. def get_system_unit_start_mqtt_topic(self, *, unit_name: str) -> str:
  86. return self._mqtt_topic_prefix + "/unit/system/" + unit_name + "/start"
  87. def get_system_unit_stop_mqtt_topic(self, *, unit_name: str) -> str:
  88. return self._mqtt_topic_prefix + "/unit/system/" + unit_name + "/stop"
  89. def get_system_unit_restart_mqtt_topic(self, *, unit_name: str) -> str:
  90. return self._mqtt_topic_prefix + "/unit/system/" + unit_name + "/restart"
  91. @property
  92. def monitored_system_unit_names(self) -> typing.List[str]:
  93. return self._monitored_system_unit_names
  94. @property
  95. def controlled_system_unit_names(self) -> typing.List[str]:
  96. return self._controlled_system_unit_names
  97. @property
  98. def shutdown_lock_acquired(self) -> bool:
  99. return self._shutdown_lock is not None
  100. def acquire_shutdown_lock(self) -> None:
  101. with self._shutdown_lock_mutex:
  102. assert self._shutdown_lock is None
  103. # https://www.freedesktop.org/wiki/Software/systemd/inhibit/
  104. (self._shutdown_lock,) = self._login_manager.Inhibit(
  105. what="shutdown",
  106. who="systemctl-mqtt",
  107. why="Report shutdown via MQTT",
  108. mode="delay",
  109. )
  110. assert isinstance(
  111. self._shutdown_lock, jeepney.fds.FileDescriptor
  112. ), self._shutdown_lock
  113. _LOGGER.debug("acquired shutdown inhibitor lock")
  114. def release_shutdown_lock(self) -> None:
  115. with self._shutdown_lock_mutex:
  116. if self._shutdown_lock:
  117. self._shutdown_lock.close()
  118. _LOGGER.debug("released shutdown inhibitor lock")
  119. self._shutdown_lock = None
  120. @property
  121. def _preparing_for_shutdown_topic(self) -> str:
  122. return self.mqtt_topic_prefix + "/preparing-for-shutdown"
  123. async def _publish_preparing_for_shutdown(
  124. self, *, mqtt_client: aiomqtt.Client, active: bool
  125. ) -> None:
  126. topic = self._preparing_for_shutdown_topic
  127. # pylint: disable=protected-access
  128. payload = systemctl_mqtt._mqtt.encode_bool(active)
  129. _LOGGER.info("publishing %r on %s", payload, topic)
  130. await mqtt_client.publish(topic=topic, payload=payload, retain=False)
  131. async def preparing_for_shutdown_handler(
  132. self, active: bool, mqtt_client: aiomqtt.Client
  133. ) -> None:
  134. active = bool(active)
  135. await self._publish_preparing_for_shutdown(
  136. mqtt_client=mqtt_client, active=active
  137. )
  138. if active:
  139. self.release_shutdown_lock()
  140. else:
  141. self.acquire_shutdown_lock()
  142. async def publish_preparing_for_shutdown(self, mqtt_client: aiomqtt.Client) -> None:
  143. try:
  144. ((return_type, active),) = self._login_manager.Get("PreparingForShutdown")
  145. except jeepney.wrappers.DBusErrorResponse as exc:
  146. _LOGGER.error(
  147. "failed to read logind's PreparingForShutdown property: %s", exc
  148. )
  149. return
  150. assert return_type == "b", return_type
  151. assert isinstance(active, bool), active
  152. await self._publish_preparing_for_shutdown(
  153. mqtt_client=mqtt_client, active=active
  154. )
  155. async def publish_homeassistant_device_config(
  156. self, mqtt_client: aiomqtt.Client
  157. ) -> None:
  158. # <discovery_prefix>/<component>/[<node_id>/]<object_id>/config
  159. # https://www.home-assistant.io/integrations/mqtt/#mqtt-discovery
  160. discovery_topic = "/".join(
  161. (
  162. self._homeassistant_discovery_prefix,
  163. "device",
  164. self._homeassistant_discovery_object_id,
  165. "config",
  166. )
  167. )
  168. hostname = (
  169. # pylint: disable=protected-access; function in internal module
  170. systemctl_mqtt._utils.get_hostname()
  171. )
  172. package_metadata = importlib.metadata.metadata(__name__)
  173. unique_id_prefix = "systemctl-mqtt-" + hostname
  174. config = {
  175. "device": {"identifiers": [hostname], "name": hostname},
  176. "origin": {
  177. "name": package_metadata["Name"],
  178. "sw_version": package_metadata["Version"],
  179. "support_url": package_metadata["Home-page"],
  180. },
  181. "availability": {"topic": self.mqtt_availability_topic},
  182. "components": {
  183. "logind/preparing-for-shutdown": {
  184. "unique_id": unique_id_prefix + "-logind-preparing-for-shutdown",
  185. "object_id": f"{hostname}_logind_preparing_for_shutdown", # entity id
  186. "name": "preparing for shutdown", # home assistant prepends device name
  187. "platform": "binary_sensor",
  188. "state_topic": self._preparing_for_shutdown_topic,
  189. # pylint: disable=protected-access
  190. "payload_on": systemctl_mqtt._mqtt.encode_bool(True),
  191. "payload_off": systemctl_mqtt._mqtt.encode_bool(False),
  192. },
  193. },
  194. }
  195. for mqtt_topic_suffix in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.keys():
  196. # false positive warning by mypy:
  197. # > Unsupported target for indexed assignment
  198. config["components"]["logind/" + mqtt_topic_suffix] = { # type: ignore
  199. "unique_id": unique_id_prefix + "-logind-" + mqtt_topic_suffix,
  200. "object_id": hostname
  201. + "_logind_"
  202. + mqtt_topic_suffix.replace("-", "_"), # entity id
  203. "name": mqtt_topic_suffix.replace("-", " "),
  204. "platform": "button",
  205. "command_topic": self.mqtt_topic_prefix + "/" + mqtt_topic_suffix,
  206. }
  207. for unit_name in self._monitored_system_unit_names:
  208. config["components"]["unit/system/" + unit_name + "/active-state"] = { # type: ignore
  209. "unique_id": f"{unique_id_prefix}-unit-system-{unit_name}-active-state",
  210. "object_id": f"{hostname}_unit_system_{unit_name}_active_state",
  211. "name": f"{unit_name} active state",
  212. "platform": "sensor",
  213. "state_topic": self.get_system_unit_active_state_mqtt_topic(
  214. unit_name=unit_name
  215. ),
  216. }
  217. for unit_name in self._controlled_system_unit_names:
  218. config["components"]["unit/system/" + unit_name + "/start"] = { # type: ignore
  219. "unique_id": f"{unique_id_prefix}-unit-system-{unit_name}-start",
  220. "object_id": f"{hostname}_unit_system_{unit_name}_start",
  221. "name": f"{unit_name} start",
  222. "platform": "button",
  223. "command_topic": self.get_system_unit_start_mqtt_topic(
  224. unit_name=unit_name
  225. ),
  226. }
  227. config["components"]["unit/system/" + unit_name + "/stop"] = { # type: ignore
  228. "unique_id": f"{unique_id_prefix}-unit-system-{unit_name}-stop",
  229. "object_id": f"{hostname}_unit_system_{unit_name}_stop",
  230. "name": f"{unit_name} stop",
  231. "platform": "button",
  232. "command_topic": self.get_system_unit_stop_mqtt_topic(
  233. unit_name=unit_name
  234. ),
  235. }
  236. config["components"]["unit/system/" + unit_name + "/restart"] = { # type: ignore
  237. "unique_id": f"{unique_id_prefix}-unit-system-{unit_name}-restart",
  238. "object_id": f"{hostname}_unit_system_{unit_name}_restart",
  239. "name": f"{unit_name} restart",
  240. "platform": "button",
  241. "command_topic": self.get_system_unit_restart_mqtt_topic(
  242. unit_name=unit_name
  243. ),
  244. }
  245. _LOGGER.debug("publishing home assistant config on %s", discovery_topic)
  246. await mqtt_client.publish(
  247. topic=discovery_topic, payload=json.dumps(config), retain=False
  248. )
  249. class _MQTTAction(metaclass=abc.ABCMeta):
  250. @abc.abstractmethod
  251. def trigger(self, state: _State) -> None:
  252. pass # pragma: no cover
  253. def __str__(self) -> str:
  254. return type(self).__name__
  255. class _MQTTActionSchedulePoweroff(_MQTTAction):
  256. # pylint: disable=too-few-public-methods
  257. def trigger(self, state: _State) -> None:
  258. # pylint: disable=protected-access
  259. systemctl_mqtt._dbus.login_manager.schedule_shutdown(
  260. action="poweroff", delay=state.poweroff_delay
  261. )
  262. class _MQTTActionStartUnit(_MQTTAction):
  263. # pylint: disable=protected-access,too-few-public-methods
  264. def __init__(self, unit_name: str):
  265. self._unit_name = unit_name
  266. def trigger(self, state: _State) -> None:
  267. systemctl_mqtt._dbus.service_manager.start_unit(unit_name=self._unit_name)
  268. class _MQTTActionStopUnit(_MQTTAction):
  269. # pylint: disable=protected-access,too-few-public-methods
  270. def __init__(self, unit_name: str):
  271. self._unit_name = unit_name
  272. def trigger(self, state: _State) -> None:
  273. systemctl_mqtt._dbus.service_manager.stop_unit(unit_name=self._unit_name)
  274. class _MQTTActionRestartUnit(_MQTTAction):
  275. # pylint: disable=protected-access,too-few-public-methods
  276. def __init__(self, unit_name: str):
  277. self._unit_name = unit_name
  278. def trigger(self, state: _State) -> None:
  279. systemctl_mqtt._dbus.service_manager.restart_unit(unit_name=self._unit_name)
  280. class _MQTTActionLockAllSessions(_MQTTAction):
  281. # pylint: disable=too-few-public-methods
  282. def trigger(self, state: _State) -> None:
  283. # pylint: disable=protected-access
  284. systemctl_mqtt._dbus.login_manager.lock_all_sessions()
  285. class _MQTTActionSuspend(_MQTTAction):
  286. # pylint: disable=too-few-public-methods
  287. def trigger(self, state: _State) -> None:
  288. # pylint: disable=protected-access
  289. systemctl_mqtt._dbus.login_manager.suspend()
  290. _MQTT_TOPIC_SUFFIX_ACTION_MAPPING = {
  291. "poweroff": _MQTTActionSchedulePoweroff(),
  292. "lock-all-sessions": _MQTTActionLockAllSessions(),
  293. "suspend": _MQTTActionSuspend(),
  294. }
  295. async def _mqtt_message_loop(*, state: _State, mqtt_client: aiomqtt.Client) -> None:
  296. action_by_topic: typing.Dict[str, _MQTTAction] = {}
  297. for topic_suffix, action in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.items():
  298. topic = state.mqtt_topic_prefix + "/" + topic_suffix
  299. _LOGGER.info("subscribing to %s", topic)
  300. await mqtt_client.subscribe(topic)
  301. action_by_topic[topic] = action
  302. for unit_name in state.controlled_system_unit_names:
  303. for topic_suffix, action_class in [
  304. ("start", _MQTTActionStartUnit),
  305. ("stop", _MQTTActionStopUnit),
  306. ("restart", _MQTTActionRestartUnit),
  307. ]:
  308. topic = (
  309. state.mqtt_topic_prefix
  310. + "/unit/system/"
  311. + unit_name
  312. + "/"
  313. + topic_suffix
  314. )
  315. _LOGGER.info("subscribing to %s", topic)
  316. await mqtt_client.subscribe(topic)
  317. action_by_topic[topic] = action_class(unit_name=unit_name)
  318. async for message in mqtt_client.messages:
  319. if message.retain:
  320. _LOGGER.info("ignoring retained message on topic %r", message.topic.value)
  321. else:
  322. _LOGGER.debug(
  323. "received message on topic %r: %r", message.topic.value, message.payload
  324. )
  325. action_by_topic[message.topic.value].trigger(state=state)
  326. async def _dbus_signal_loop_preparing_for_shutdown(
  327. *,
  328. state: _State,
  329. mqtt_client: aiomqtt.Client,
  330. dbus_router: jeepney.io.asyncio.DBusRouter,
  331. bus_proxy: jeepney.io.asyncio.Proxy,
  332. ) -> None:
  333. preparing_for_shutdown_match_rule = (
  334. # pylint: disable=protected-access
  335. systemctl_mqtt._dbus.login_manager.get_login_manager_signal_match_rule(
  336. "PrepareForShutdown"
  337. )
  338. )
  339. assert await bus_proxy.AddMatch(preparing_for_shutdown_match_rule) == ()
  340. with dbus_router.filter(preparing_for_shutdown_match_rule) as queue:
  341. while True:
  342. message: jeepney.low_level.Message = await queue.get()
  343. (preparing_for_shutdown,) = message.body
  344. await state.preparing_for_shutdown_handler(
  345. active=preparing_for_shutdown, mqtt_client=mqtt_client
  346. )
  347. queue.task_done()
  348. async def _get_unit_path(
  349. *, service_manager: jeepney.io.asyncio.Proxy, unit_name: str
  350. ) -> str:
  351. (path,) = await service_manager.GetUnit(name=unit_name)
  352. return path
  353. async def _dbus_signal_loop_unit( # pylint: disable=too-many-arguments
  354. *,
  355. state: _State,
  356. mqtt_client: aiomqtt.Client,
  357. dbus_router: jeepney.io.asyncio.DBusRouter,
  358. bus_proxy: jeepney.io.asyncio.Proxy,
  359. unit_name: str,
  360. unit_path: str,
  361. ) -> None:
  362. unit_proxy = jeepney.io.asyncio.Proxy(
  363. # pylint: disable=protected-access
  364. msggen=systemctl_mqtt._dbus.service_manager.Unit(object_path=unit_path),
  365. router=dbus_router,
  366. )
  367. unit_properties_changed_match_rule = jeepney.MatchRule(
  368. type="signal",
  369. interface="org.freedesktop.DBus.Properties",
  370. member="PropertiesChanged",
  371. path=unit_path,
  372. )
  373. assert (await bus_proxy.AddMatch(unit_properties_changed_match_rule)) == ()
  374. # > Table 1. Unit ACTIVE states …
  375. # > active Started, bound, plugged in, …
  376. # > inactive Stopped, unbound, unplugged, …
  377. # > failed … process returned error code on exit, crashed, an operation
  378. # . timed out, or after too many restarts).
  379. # > activating Changing from inactive to active.
  380. # > deactivating Changing from active to inactive.
  381. # > maintenance Unit is inactive and … maintenance … in progress.
  382. # > reloading Unit is active and it is reloading its configuration.
  383. # > refreshing Unit is active and a new mount is being activated in its
  384. # . namespace.
  385. # https://web.archive.org/web/20250101121304/https://www.freedesktop.org/software/systemd/man/latest/org.freedesktop.systemd1.html
  386. active_state_topic = state.get_system_unit_active_state_mqtt_topic(
  387. unit_name=unit_name
  388. )
  389. ((_, last_active_state),) = await unit_proxy.Get(property_name="ActiveState")
  390. await mqtt_client.publish(topic=active_state_topic, payload=last_active_state)
  391. with dbus_router.filter(unit_properties_changed_match_rule) as queue:
  392. while True:
  393. await queue.get()
  394. ((_, current_active_state),) = await unit_proxy.Get(
  395. property_name="ActiveState"
  396. )
  397. if current_active_state != last_active_state:
  398. await mqtt_client.publish(
  399. topic=active_state_topic, payload=current_active_state
  400. )
  401. last_active_state = current_active_state
  402. queue.task_done()
  403. async def _dbus_signal_loop(*, state: _State, mqtt_client: aiomqtt.Client) -> None:
  404. async with jeepney.io.asyncio.open_dbus_router(bus="SYSTEM") as router:
  405. # router: jeepney.io.asyncio.DBusRouter
  406. bus_proxy = jeepney.io.asyncio.Proxy(
  407. msggen=jeepney.bus_messages.message_bus, router=router
  408. )
  409. system_service_manager = jeepney.io.asyncio.Proxy(
  410. # pylint: disable=protected-access
  411. msggen=systemctl_mqtt._dbus.service_manager.ServiceManager(),
  412. router=router,
  413. )
  414. await asyncio.gather(
  415. *[
  416. _dbus_signal_loop_preparing_for_shutdown(
  417. state=state,
  418. mqtt_client=mqtt_client,
  419. dbus_router=router,
  420. bus_proxy=bus_proxy,
  421. )
  422. ]
  423. + [
  424. _dbus_signal_loop_unit(
  425. state=state,
  426. mqtt_client=mqtt_client,
  427. dbus_router=router,
  428. bus_proxy=bus_proxy,
  429. unit_name=unit_name,
  430. unit_path=await _get_unit_path(
  431. service_manager=system_service_manager, unit_name=unit_name
  432. ),
  433. )
  434. for unit_name in state.monitored_system_unit_names
  435. ],
  436. return_exceptions=False,
  437. )
  438. async def _run( # pylint: disable=too-many-arguments
  439. *,
  440. mqtt_host: str,
  441. mqtt_port: int,
  442. mqtt_username: typing.Optional[str],
  443. mqtt_password: typing.Optional[str],
  444. mqtt_topic_prefix: str,
  445. homeassistant_discovery_prefix: str,
  446. homeassistant_discovery_object_id: str,
  447. poweroff_delay: datetime.timedelta,
  448. monitored_system_unit_names: typing.List[str],
  449. controlled_system_unit_names: typing.List[str],
  450. mqtt_disable_tls: bool = False,
  451. ) -> None:
  452. state = _State(
  453. mqtt_topic_prefix=mqtt_topic_prefix,
  454. homeassistant_discovery_prefix=homeassistant_discovery_prefix,
  455. homeassistant_discovery_object_id=homeassistant_discovery_object_id,
  456. poweroff_delay=poweroff_delay,
  457. monitored_system_unit_names=monitored_system_unit_names,
  458. controlled_system_unit_names=controlled_system_unit_names,
  459. )
  460. _LOGGER.info(
  461. "connecting to MQTT broker %s:%d (TLS %s)",
  462. mqtt_host,
  463. mqtt_port,
  464. "disabled" if mqtt_disable_tls else "enabled",
  465. )
  466. if mqtt_password and not mqtt_username:
  467. raise ValueError("Missing MQTT username")
  468. async with aiomqtt.Client( # raises aiomqtt.MqttError
  469. hostname=mqtt_host,
  470. port=mqtt_port,
  471. # > The settings [...] usually represent a higher security level than
  472. # > when calling the SSLContext constructor directly.
  473. # https://web.archive.org/web/20230714183106/https://docs.python.org/3/library/ssl.html
  474. tls_context=None if mqtt_disable_tls else ssl.create_default_context(),
  475. username=None if mqtt_username is None else mqtt_username,
  476. password=None if mqtt_password is None else mqtt_password,
  477. will=aiomqtt.Will( # e.g. on SIGTERM & SIGKILL
  478. topic=state.mqtt_availability_topic,
  479. payload=_MQTT_PAYLOAD_NOT_AVAILABLE,
  480. retain=True,
  481. ),
  482. ) as mqtt_client:
  483. _LOGGER.debug("connected to MQTT broker %s:%d", mqtt_host, mqtt_port)
  484. if not state.shutdown_lock_acquired:
  485. state.acquire_shutdown_lock()
  486. await state.publish_homeassistant_device_config(mqtt_client=mqtt_client)
  487. await state.publish_preparing_for_shutdown(mqtt_client=mqtt_client)
  488. try:
  489. await mqtt_client.publish(
  490. topic=state.mqtt_availability_topic,
  491. payload=_MQTT_PAYLOAD_AVAILABLE,
  492. retain=True,
  493. )
  494. # asynpio.TaskGroup added in python3.11
  495. await asyncio.gather(
  496. _mqtt_message_loop(state=state, mqtt_client=mqtt_client),
  497. _dbus_signal_loop(state=state, mqtt_client=mqtt_client),
  498. return_exceptions=False,
  499. )
  500. finally: # e.g. on SIGINT
  501. # https://web.archive.org/web/20250101080719/https://github.com/empicano/aiomqtt/issues/28
  502. await mqtt_client.publish(
  503. topic=state.mqtt_availability_topic,
  504. payload=_MQTT_PAYLOAD_NOT_AVAILABLE,
  505. retain=True,
  506. )
  507. def _main() -> None:
  508. logging.basicConfig(
  509. level=logging.INFO,
  510. format="%(asctime)s:%(levelname)s:%(message)s",
  511. datefmt="%Y-%m-%dT%H:%M:%S%z",
  512. )
  513. argparser = argparse.ArgumentParser(
  514. description="MQTT client triggering & reporting shutdown on systemd-based systems",
  515. )
  516. argparser.add_argument(
  517. "--log-level",
  518. choices=_ARGUMENT_LOG_LEVEL_MAPPING.keys(),
  519. default="info",
  520. help="log level (default: %(default)s)",
  521. )
  522. argparser.add_argument("--mqtt-host", type=str, required=True)
  523. argparser.add_argument(
  524. "--mqtt-port",
  525. type=int,
  526. help=f"default {_MQTT_DEFAULT_TLS_PORT} ({_MQTT_DEFAULT_PORT} with --mqtt-disable-tls)",
  527. )
  528. argparser.add_argument("--mqtt-username", type=str)
  529. argparser.add_argument("--mqtt-disable-tls", action="store_true")
  530. password_argument_group = argparser.add_mutually_exclusive_group()
  531. password_argument_group.add_argument("--mqtt-password", type=str)
  532. password_argument_group.add_argument(
  533. "--mqtt-password-file",
  534. type=pathlib.Path,
  535. metavar="PATH",
  536. dest="mqtt_password_path",
  537. help="stripping trailing newline",
  538. )
  539. argparser.add_argument(
  540. "--mqtt-topic-prefix",
  541. type=str,
  542. # pylint: disable=protected-access
  543. default="systemctl/" + systemctl_mqtt._utils.get_hostname(),
  544. help="default: %(default)s",
  545. )
  546. # https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix
  547. argparser.add_argument(
  548. "--homeassistant-discovery-prefix",
  549. type=str,
  550. default="homeassistant",
  551. help="home assistant's prefix for discovery topics" + " (default: %(default)s)",
  552. )
  553. argparser.add_argument(
  554. "--homeassistant-discovery-object-id",
  555. type=str,
  556. # pylint: disable=protected-access
  557. default=systemctl_mqtt._homeassistant.get_default_discovery_object_id(),
  558. help="part of discovery topic (default: %(default)s)",
  559. )
  560. argparser.add_argument(
  561. "--poweroff-delay-seconds", type=float, default=4.0, help="default: %(default)s"
  562. )
  563. argparser.add_argument(
  564. "--monitor-system-unit",
  565. type=str,
  566. metavar="UNIT_NAME",
  567. dest="monitored_system_unit_names",
  568. action="append",
  569. help="e.g. --monitor-system-unit ssh.service --monitor-system-unit custom.service",
  570. )
  571. argparser.add_argument(
  572. "--control-system-unit",
  573. type=str,
  574. metavar="UNIT_NAME",
  575. dest="controlled_system_unit_names",
  576. action="append",
  577. help="e.g. --control-system-unit ansible-pull.service --control-system-unit custom.service",
  578. )
  579. args = argparser.parse_args()
  580. logging.root.setLevel(_ARGUMENT_LOG_LEVEL_MAPPING[args.log_level])
  581. if args.mqtt_port:
  582. mqtt_port = args.mqtt_port
  583. elif args.mqtt_disable_tls:
  584. mqtt_port = _MQTT_DEFAULT_PORT
  585. else:
  586. mqtt_port = _MQTT_DEFAULT_TLS_PORT
  587. if args.mqtt_password_path:
  588. # .read_text() replaces \r\n with \n
  589. mqtt_password = args.mqtt_password_path.read_bytes().decode()
  590. if mqtt_password.endswith("\r\n"):
  591. mqtt_password = mqtt_password[:-2]
  592. elif mqtt_password.endswith("\n"):
  593. mqtt_password = mqtt_password[:-1]
  594. else:
  595. mqtt_password = args.mqtt_password
  596. # pylint: disable=protected-access
  597. if not systemctl_mqtt._homeassistant.validate_discovery_object_id(
  598. args.homeassistant_discovery_object_id
  599. ):
  600. raise ValueError(
  601. # pylint: disable=protected-access
  602. "invalid home assistant discovery object id"
  603. f" {args.homeassistant_discovery_object_id!r} (length >= 1"
  604. ", allowed characters:"
  605. f" {systemctl_mqtt._homeassistant.NODE_ID_ALLOWED_CHARS})"
  606. "\nchange --homeassistant-discovery-object-id"
  607. )
  608. asyncio.run(
  609. _run(
  610. mqtt_host=args.mqtt_host,
  611. mqtt_port=mqtt_port,
  612. mqtt_disable_tls=args.mqtt_disable_tls,
  613. mqtt_username=args.mqtt_username,
  614. mqtt_password=mqtt_password,
  615. mqtt_topic_prefix=args.mqtt_topic_prefix,
  616. homeassistant_discovery_prefix=args.homeassistant_discovery_prefix,
  617. homeassistant_discovery_object_id=args.homeassistant_discovery_object_id,
  618. poweroff_delay=datetime.timedelta(seconds=args.poweroff_delay_seconds),
  619. monitored_system_unit_names=args.monitored_system_unit_names or [],
  620. controlled_system_unit_names=args.controlled_system_unit_names or [],
  621. )
  622. )