__init__.py 27 KB

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