#!/usr/bin/env python
# -*- encoding: iso-8859-1 -*-
#
# apt_ubuntu
#
# Plugin to monitor packages that should be installed on Ubuntu systems.
#
# Author: Stefan Daniel Schwarz <munin@wolfram.ravenwolf.de>
#
# v1.0 2008-11-07 - First draft
# v1.1 2008-11-08 - critical = #: First # critical, rest warning
# v1.2 2008-11-09 - Code cleanup for MuninExchange submission
#
# Usage: place in /etc/munin/plugins/ (or link it there using ln -s)
#
# Parameters understood:
#
#       config   (required)
#       autoconf (optional - used by munin-config)
#
# Magic markers - optional - used by installation scripts and
# munin-config:
#
#  #%# capabilities=autoconf
#  #%# family=contrib

import os
import sys
import warnings

###########################################################
title = 'Upgradable packages'  # 'Upgradeable packages'
vlabel = 'Total packages'
other = 'other'
total = 'total'

archives = ['security', 'updates', 'proposed', 'backports']
colour = ['ff0000', '22ff22', '0022ff', '00aaaa', 'ff00ff']
origins = ['Ubuntu']

critical = 1
###########################################################

warnings.filterwarnings('ignore', 'apt API not stable yet', FutureWarning)


def autoconf():
    if os.path.exists('/etc/lsb-release'):
        for line in open('/etc/lsb-release'):
            if line.strip() == 'DISTRIB_ID=Ubuntu':
                try:
                    import apt
                except ImportError:
                    print('no (python-apt not installed)')
                    sys.exit(0)
                cache = apt.Cache()
                if 'update-notifier-common' not in cache:
                    print('no (update-notifier-common not found)')
                    sys.exit(0)
                if not cache['update-notifier-common'].isInstalled:
                    print('no (update-notifier-common not installed)')
                    sys.exit(0)
                if not os.path.exists('/etc/apt/apt.conf.d/10periodic'):
                    print('no (/etc/apt/apt.conf.d/10periodic not found)')
                    sys.exit(1)
                for line in open('/etc/apt/apt.conf.d/10periodic'):
                    if line.strip() == 'APT::Periodic::Update-Package-Lists "1";':
                        print('yes')
                        sys.exit(0)
                print('no (APT::Periodic::Update-Package-Lists not "1")')
                sys.exit(0)
    print('no (missing /etc/lsb-release file)')
    sys.exit(0)


def config():
    print('graph_category security')
    print('graph_title %s' % (title))
    print('graph_vlabel %s' % (vlabel))
    for i, archive in enumerate(archives + [other]):
        if len(colour) > i:
            print('%s.colour %s' % (archive, colour[i]))
        if i < critical:
            print('%s.critical 0:0' % (archive))
        if i == 0:
            print('%s.draw AREA' % (archive))
        else:
            print('%s.draw STACK' % (archive))
        print('%s.label %s' % (archive, archive))
        if i + 1 > critical:
            print('%s.warning 0:0' % (archive))
    print('total.colour 000000')
    print('total.draw LINE1')
    print('total.label %s' % (total))
    sys.exit(0)


def check_origin(pkg):
    if pkg.candidate.origins:
        for archive in archives:
            for origin in pkg.candidate.origins:
                a = origin.archive.split('-')[origin.archive.count('-')]
                if (a == archive) and (origin.origin in origins):
                    return a
    return other


if len(sys.argv) > 1:
    if sys.argv[1] == 'autoconf':
        autoconf()
    elif sys.argv[1] == 'config':
        config()
    elif sys.argv[1]:
        print('unknown argument "' + sys.argv[1] + '"')
        sys.exit(1)

try:
    import apt
    import apt_pkg
except ImportError:
    print("The module 'apt' is currently not installed.  You can install it by typing:\n"
          "sudo apt-get install python-apt\nImportError: No module named apt")
    sys.exit(1)


pkgs = {}
total = 0
for pkg in apt.Cache():
    if (pkg.is_upgradable) and (pkg._pkg.selected_state != apt_pkg.SELSTATE_HOLD):
        a = check_origin(pkg)
        pkgs[a] = pkgs.get(a, 0) + 1
        total += 1

for archive in archives + [other]:
    print('%s.value %s' % (archive, pkgs.pop(archive, 0)))

print('total.value %s' % (total))