#!/usr/bin/env python
# collectd-systemd - systemd.py :
#     A collectd plugin to read systemd unit and job counts
#
# Copyright 2017 Luke Shumaker <lukeshu@lukeshu.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# This should be both valid Python 2 and 3, as the collectd Python
# plugin can be either Python 2 or 3 (decided at compile time).
import sys
import dbus

# For compatibility with Python 2 & 3
if sys.version_info.major is 2:
    def iteritems(d):
        return d.iteritems()
else:
    def iteritems(d):
        return d.items()

# DBus utility function
def get_property(dbus_iface, propname):
    return dbus_iface.proxy_object.Get(dbus_iface.dbus_interface, propname, dbus_interface='org.freedesktop.DBus.Properties')

# The main program

def init():
    global system_manager
    system_manager = dbus.Interface(dbus.SystemBus().get_object("org.freedesktop.systemd1", "/org/freedesktop/systemd1"), "org.freedesktop.systemd1.Manager")

def read():
    global system_manager

    dat = {
        # These are the values of the
        # basic/unit-name.h:UnitActiveStates enum.
        ('sd_units', 'active'): 0,
        ('sd_units', 'reloading'): 0,
        ('sd_units', 'inactive'): 0,
        ('sd_units', 'failed'): 0,
        ('sd_units', 'activating'): 0,
        ('sd_units', 'deactivating'): 0,
    }
    for unit in system_manager.ListUnits():
        key = ('sd_units', str(unit[3]))
        dat[key] = dat[key] + 1

    dat[('sd_jobs', 'queued')]    = int(get_property(system_manager, 'NJobs'))
    dat[('sd_jobs', 'installed')] = int(get_property(system_manager, 'NInstalledJobs'))
    dat[('sd_jobs', 'failed')]    = int(get_property(system_manager, 'NFailedJobs'))

    return dat

if __name__ == '__main__':
    init()
    for (key, val) in iteritems(read()):
        print("%s => %s" % (key, val))
else:
    import collectd
    def _config(conf):
        pass
    def _read():
        for (key, val) in iteritems(read()):
            collectd.Values(
                plugin='systemd',
                type=key[0],
                type_instance=key[1]
            ).dispatch(
                values=[val]
            )
    collectd.register_config(_config)
    collectd.register_init(init)
    collectd.register_read(_read, 1)