1
0

__init__.py 7.7 KB

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