# Source code for nobodd.systemd

    # nobodd: a boot configuration tool for the Raspberry Pi
    #
    # Copyright (c) 2024 Dave Jones <dave.jones@canonical.com>
    # Copyright (c) 2024 Canonical Ltd.
    #
    # SPDX-License-Identifier: GPL-3.0

    """
    Provides a simple interface to systemd's notification and watchdog services.

    .. autoclass:: Systemd
    """

    import os
    import socket

    from . import lang

    class Systemd:
        """
        Provides a simple interface to systemd's notification and watchdog
        services. It is suggested applications obtain a single, top-level instance
        of this class via :func:`get_systemd` and use it to communicate with
        systemd.
        """
        __slots__ = ('_socket',)
        LISTEN_FDS_START = 3

        def __init__(self, address=None):
            # Remove NOTIFY_SOCKET implicitly so child processes don't inherit it
            self._socket = None
            if address is None:
                address = os.environ.pop('NOTIFY_SOCKET', None)
            if address is not None:
                if len(address) <= 1 or address[0] not in ('@', '/'):
                    return None
                if address[0] == '@':
                    address = '\0' + address[1:] # abstract namespace socket
                self._socket = socket.socket(
                    socket.AF_UNIX, socket.SOCK_DGRAM | socket.SOCK_CLOEXEC)
                try:
                    self._socket.connect(address)
                except IOError:
                    self._socket = None

        def available(self):
            """
            If systemd's notification socket is not available, raises
            :exc:`RuntimeError`. Services expecting systemd notifications to be
            available can call this to assert that notifications will be noticed.
            """
            if self._socket is None:
                raise RuntimeError(lang._(
                    'systemd notification socket unavailable'))

        def notify(self, state):
            """
            Send a notification to systemd. *state* is a string type (if it is a
            unicode string it will be encoded with the 'ascii' codec).
            """
            if self._socket is not None:
                if isinstance(state, str):
                    state = state.encode('ascii')
                self._socket.sendall(state)

        def ready(self):
            """
            Notify systemd that service startup is complete.
            """
            self.notify('READY=1')

        def reloading(self):
            """
            Notify systemd that the service is reloading its configuration. Call
            :func:`ready` when reload is complete.
            """
            self.notify('RELOADING=1')

        def stopping(self):
            """
            Notify systemd that the service is stopping.
            """
            self.notify('STOPPING=1')

        def extend_timeout(self, timeout):
            """
            Notify systemd to extend the start / stop timeout by *timeout* seconds.
            A timeout will occur if the service does not call :func:`ready` or
            terminate within *timeout* seconds but *only* if the original timeout
            (set in the systemd configuration) has already been exceeded.

            For example, if the stopping timeout is configured as 90s, and the
            service calls :func:`stopping`, systemd expects the service to
            terminate within 90s. After 10s the service calls
            :func:`extend_timeout` with a *timeout* of 10s. 20s later the service
            has not yet terminated but systemd does *not* consider the timeout
            expired as only 30s have elapsed of the original 90s timeout.
            """
            self.notify(f'EXTEND_TIMEOUT_USEC={timeout * 1000000:d}')

        def watchdog_ping(self):
            """
            Ping the systemd watchdog. This must be done periodically if
            :func:`watchdog_period` returns a value other than ``None``.
            """
            self.notify('WATCHDOG=1')

        def watchdog_reset(self, timeout):
            """
            Reset the systemd watchdog timer to *timeout* seconds.
            """
            self.notify(f'WATCHDOG_USEC={timeout * 1000000:d}')

        def watchdog_period(self):
            """
            Returns the time (in seconds) before which systemd expects the process
            to call :func:`watchdog_ping`. If a watchdog timeout is not set, the
            function returns ``None``.
            """
            timeout = os.environ.get('WATCHDOG_USEC')
            if timeout is not None:
                pid = os.environ.get('WATCHDOG_PID')
                if pid is None or int(pid) == os.getpid():
                    return int(timeout) / 1000000
            return None

        def watchdog_clean(self):
            """
            Unsets the watchdog environment variables so that no future child
            processes will inherit them.

            .. warning::

                After calling this function, :func:`watchdog_period` will return
                ``None`` but systemd will continue expecting :func:`watchdog_ping`
                to be called periodically. In other words, you should call
                :func:`watchdog_period` first and store its result somewhere before
                calling this function.
            """
            os.environ.pop('WATCHDOG_USEC', None)
            os.environ.pop('WATCHDOG_PID', None)

        def main_pid(self, pid=None):
            """
            Report the main PID of the process to systemd (for services that
            confuse systemd with their forking behaviour). If *pid* is None,
            :func:`os.getpid` is called to determine the calling process' PID.
            """
            if pid is None:
                pid = os.getpid()
            self.notify(f'MAINPID={pid:d}')

        def listen_fds(self):
            """
            Return file-descriptors passed to the service by systemd, e.g. as part
            of socket activation or file descriptor stores. It returns a
            :class:`dict` mapping each file-descriptor to its name, or the string
            "unknown" if no name was given.
            """
            try:
                if int(os.environ['LISTEN_PID']) != os.getpid():
                    raise ValueError(lang._('wrong LISTEN_PID'))
                fds = int(os.environ['LISTEN_FDS'])
            except (ValueError, KeyError):
                return {}
            try:
                names = os.environ['LISTEN_FDNAMES'].split(':')
            except KeyError:
                names = ['unknown'] * fds
            if len(names) != fds:
                return {}
            return {
                fd: name
                for fd, name in zip(
                    range(self.LISTEN_FDS_START, self.LISTEN_FDS_START + fds),
                    names)
            }

    _SYSTEMD = None

    def get_systemd():
        """
        Return a single top-level instance of :class:`Systemd`; repeated calls will
        return the same instance.
        """
        global _SYSTEMD
        if _SYSTEMD is None:
            _SYSTEMD = Systemd()
        return _SYSTEMD