__init__.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. # systemctl-mqtt - MQTT client triggering 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 argparse
  18. import datetime
  19. import functools
  20. import logging
  21. import pathlib
  22. import socket
  23. import typing
  24. import dbus
  25. import paho.mqtt.client
  26. # black keeps inserting a blank line above
  27. # https://pygobject.readthedocs.io/en/latest/getting_started.html#ubuntu-logo-ubuntu-debian-logo-debian
  28. import gi.repository.GLib # pylint-import-requirements: imports=PyGObject
  29. _LOGGER = logging.getLogger(__name__)
  30. _SHUTDOWN_DELAY = datetime.timedelta(seconds=4)
  31. def _get_login_manager() -> dbus.proxies.Interface:
  32. # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html
  33. bus = dbus.SystemBus()
  34. proxy = bus.get_object(
  35. bus_name="org.freedesktop.login1", object_path="/org/freedesktop/login1"
  36. ) # type: dbus.proxies.ProxyObject
  37. # https://freedesktop.org/wiki/Software/systemd/logind/
  38. return dbus.Interface(object=proxy, dbus_interface="org.freedesktop.login1.Manager")
  39. def _schedule_shutdown(action: str) -> None:
  40. # https://github.com/systemd/systemd/blob/v237/src/systemctl/systemctl.c#L8553
  41. assert action in ["poweroff", "reboot"], action
  42. shutdown_datetime = datetime.datetime.now() + _SHUTDOWN_DELAY
  43. # datetime.datetime.isoformat(timespec=) not available in python3.5
  44. # https://github.com/python/cpython/blob/v3.5.9/Lib/datetime.py#L1552
  45. _LOGGER.info(
  46. "scheduling %s for %s", action, shutdown_datetime.strftime("%Y-%m-%d %H:%M:%S"),
  47. )
  48. shutdown_epoch_usec = int(shutdown_datetime.timestamp() * 10 ** 6)
  49. try:
  50. # $ gdbus introspect --system --dest org.freedesktop.login1 \
  51. # --object-path /org/freedesktop/login1 | grep -A 1 ScheduleShutdown
  52. # ScheduleShutdown(in s arg_0,
  53. # in t arg_1);
  54. # $ gdbus call --system --dest org.freedesktop.login1 \
  55. # --object-path /org/freedesktop/login1 \
  56. # --method org.freedesktop.login1.Manager.ScheduleShutdown \
  57. # poweroff "$(date --date=10min +%s)000000"
  58. # $ dbus-send --type=method_call --print-reply --system --dest=org.freedesktop.login1 \
  59. # /org/freedesktop/login1 \
  60. # org.freedesktop.login1.Manager.ScheduleShutdown \
  61. # string:poweroff "uint64:$(date --date=10min +%s)000000"
  62. _get_login_manager().ScheduleShutdown(action, shutdown_epoch_usec)
  63. except dbus.DBusException as exc:
  64. exc_msg = exc.get_dbus_message()
  65. if "authentication required" in exc_msg.lower():
  66. _LOGGER.error(
  67. "failed to schedule %s: unauthorized; missing polkit authorization rules?",
  68. action,
  69. )
  70. else:
  71. _LOGGER.error("failed to schedule %s: %s", action, exc_msg)
  72. class _Settings:
  73. # pylint: disable=too-few-public-methods
  74. def __init__(self, mqtt_topic_prefix: str) -> None:
  75. self._mqtt_topic_prefix = mqtt_topic_prefix
  76. @property
  77. def mqtt_topic_prefix(self) -> str:
  78. return self._mqtt_topic_prefix
  79. class _MQTTAction:
  80. # pylint: disable=too-few-public-methods
  81. def __init__(self, name: str, action: typing.Callable) -> None:
  82. self.name = name
  83. self.action = action
  84. def mqtt_message_callback(
  85. self,
  86. mqtt_client: paho.mqtt.client.Client,
  87. settings: _Settings,
  88. message: paho.mqtt.client.MQTTMessage,
  89. ) -> None:
  90. # pylint: disable=unused-argument; callback
  91. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L3416
  92. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L469
  93. _LOGGER.debug("received topic=%s payload=%r", message.topic, message.payload)
  94. if message.retain:
  95. _LOGGER.info("ignoring retained message")
  96. return
  97. _LOGGER.debug("executing action %s (%r)", self.name, self.action)
  98. self.action()
  99. _LOGGER.debug("completed action %s (%r)", self.name, self.action)
  100. _MQTT_TOPIC_SUFFIX_ACTION_MAPPING = {
  101. "poweroff": _MQTTAction(
  102. name="poweroff", action=functools.partial(_schedule_shutdown, action="poweroff")
  103. ),
  104. }
  105. def _mqtt_on_connect(
  106. mqtt_client: paho.mqtt.client.Client,
  107. settings: _Settings,
  108. flags: typing.Dict,
  109. return_code: int,
  110. ) -> None:
  111. # pylint: disable=unused-argument; callback
  112. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L441
  113. assert return_code == 0, return_code # connection accepted
  114. mqtt_broker_host, mqtt_broker_port = mqtt_client.socket().getpeername()
  115. _LOGGER.debug("connected to MQTT broker %s:%d", mqtt_broker_host, mqtt_broker_port)
  116. for topic_suffix, action in _MQTT_TOPIC_SUFFIX_ACTION_MAPPING.items():
  117. topic = settings.mqtt_topic_prefix + "/" + topic_suffix
  118. _LOGGER.info("subscribing to %s", topic)
  119. mqtt_client.subscribe(topic)
  120. mqtt_client.message_callback_add(
  121. sub=topic, callback=action.mqtt_message_callback
  122. )
  123. _LOGGER.debug(
  124. "registered MQTT callback for topic %s triggering %r", topic, action.action
  125. )
  126. def _run(
  127. mqtt_host: str,
  128. mqtt_port: int,
  129. mqtt_username: typing.Optional[str],
  130. mqtt_password: typing.Optional[str],
  131. mqtt_topic_prefix: str,
  132. ) -> None:
  133. # https://pypi.org/project/paho-mqtt/
  134. mqtt_client = paho.mqtt.client.Client(
  135. userdata=_Settings(mqtt_topic_prefix=mqtt_topic_prefix)
  136. )
  137. mqtt_client.on_connect = _mqtt_on_connect
  138. mqtt_client.tls_set(ca_certs=None) # enable tls trusting default system certs
  139. _LOGGER.info(
  140. "connecting to MQTT broker %s:%d", mqtt_host, mqtt_port,
  141. )
  142. if mqtt_username:
  143. mqtt_client.username_pw_set(username=mqtt_username, password=mqtt_password)
  144. elif mqtt_password:
  145. raise ValueError("Missing MQTT username")
  146. mqtt_client.connect(host=mqtt_host, port=mqtt_port)
  147. # loop_start runs loop_forever in a new thread (daemon)
  148. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1814
  149. # loop_forever attempts to reconnect if disconnected
  150. # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1744
  151. mqtt_client.loop_start()
  152. try:
  153. # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
  154. gi.repository.GLib.MainLoop().run()
  155. finally:
  156. # blocks until loop_forever stops
  157. _LOGGER.debug("waiting for MQTT loop to stop")
  158. mqtt_client.loop_stop()
  159. _LOGGER.debug("MQTT loop stopped")
  160. def _get_hostname() -> str:
  161. return socket.gethostname()
  162. def _main() -> None:
  163. logging.basicConfig(
  164. level=logging.DEBUG,
  165. format="%(asctime)s:%(levelname)s:%(message)s",
  166. datefmt="%Y-%m-%dT%H:%M:%S%z",
  167. )
  168. argparser = argparse.ArgumentParser(
  169. description="MQTT client triggering shutdown on systemd-based systems",
  170. formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  171. )
  172. argparser.add_argument("--mqtt-host", type=str, required=True)
  173. argparser.add_argument("--mqtt-port", type=int, default=8883)
  174. argparser.add_argument("--mqtt-username", type=str)
  175. password_argument_group = argparser.add_mutually_exclusive_group()
  176. password_argument_group.add_argument("--mqtt-password", type=str)
  177. password_argument_group.add_argument(
  178. "--mqtt-password-file",
  179. type=pathlib.Path,
  180. metavar="PATH",
  181. dest="mqtt_password_path",
  182. help="stripping trailing newline",
  183. )
  184. # https://www.home-assistant.io/docs/mqtt/discovery/#discovery_prefix
  185. argparser.add_argument(
  186. "--mqtt-topic-prefix",
  187. type=str,
  188. default="systemctl/" + _get_hostname(),
  189. help=" ", # show default
  190. )
  191. args = argparser.parse_args()
  192. if args.mqtt_password_path:
  193. # .read_text() replaces \r\n with \n
  194. mqtt_password = args.mqtt_password_path.read_bytes().decode()
  195. if mqtt_password.endswith("\r\n"):
  196. mqtt_password = mqtt_password[:-2]
  197. elif mqtt_password.endswith("\n"):
  198. mqtt_password = mqtt_password[:-1]
  199. else:
  200. mqtt_password = args.mqtt_password
  201. _run(
  202. mqtt_host=args.mqtt_host,
  203. mqtt_port=args.mqtt_port,
  204. mqtt_username=args.mqtt_username,
  205. mqtt_password=mqtt_password,
  206. mqtt_topic_prefix=args.mqtt_topic_prefix,
  207. )