__init__.py 27 KB

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