#!/usr/bin/python
#Almost completely stolen from the apt-check script,
#This script serves specfically for security updates,
#It can be used in cron to install security updates : --install-security-updates
#Can also be used as a check_mk local check with --security-update-count
#check_mk local check reference
#<0|1|2|3> <checkname> count=<current_value>;<warn>;<crit>;<min>;<max> <text output of command>



#nice apt-get -s -o Debug::NoLocking=true upgrade | grep ^Inst 

import apt_pkg
import os
import sys
from optparse import OptionParser
import gettext
import subprocess

SYNAPTIC_PINFILE = "/var/lib/synaptic/preferences"
DISTRO = subprocess.Popen(["lsb_release","-c","-s"],
                          stdout=subprocess.PIPE).communicate()[0].strip()

class OpNullProgress(object):
    def update(self, percent):
        pass
    def done(self):
        pass

def _(msg):
    return gettext.dgettext("update-notifier", msg)

def _handleException(type, value, tb):
    sys.stderr.write("E: "+ _("Unknown Error: '%s' (%s)") % (type,value))
    sys.exit(-1)

def clean(cache,depcache):
    " unmark (clean) all changes from the given depcache "
    # mvo: looping is too inefficient with the new auto-mark code
    #for pkg in cache.Packages:
    #    depcache.MarkKeep(pkg)
    depcache.init()

def saveDistUpgrade(cache,depcache):
    """ this functions mimics a upgrade but will never remove anything """
    depcache.upgrade(True)
    if depcache.del_count > 0:
        clean(cache,depcache)
    depcache.upgrade()

def isSecurityUpgrade(ver):
    " check if the given version is a security update (or masks one) "
    security_pockets = [("Ubuntu", "%s-security" % DISTRO),
                        ("gNewSense", "%s-security" % DISTRO),
                        ("Debian", "%s-updates" % DISTRO)]

    for (file, index) in ver.file_list:
        for origin, archive in security_pockets:
            if (file.archive == archive and file.origin == origin):
                return True
    return False

def write_package_names(pkgs):
	for pkg in pkgs:
		print pkg


def write_human_readable_summary(outstream, security_updates):
    " write out human summary summary to outstream "
    outstream.write(gettext.dngettext("update-notifier",
                            "%i update is a security update.",
                            "%i updates are security updates.",
                            security_updates)  % security_updates)
    outstream.write("\n")

def install_security_updates(pkgs):
	status = 0
	myenv = dict(os.environ)
	myenv['DEBIAN_FRONTEND'] = 'noninteractive'
	#print myenv
	for pkg in pkgs:
	#	print 'installing security update -> ' + pkg
		c = subprocess.Popen(["apt-get","install", "-y", pkg], env=myenv, stdout=subprocess.PIPE)
		streamdata = c.communicate()[0]
		if c.returncode != 0:
	#		print 'exit: ' + str(c.returncode) + ' ,'  + 'install failed for ' + pkg
			status = 1
	#	else:
	#		print 'exit: ' + str(c.returncode) + ' ,' + pkg + 'updated successfully.'
	return status		
		
def init():
    " init the system, be nice "
    # FIXME: do a ionice here too?
    os.nice(19)
    apt_pkg.init()
    # force apt to build its caches in memory for now to make sure
    # that there is no race when the pkgcache file gets re-generated
    apt_pkg.config.set("Dir::Cache::pkgcache","")
    
def run(options=None):

    # we are run in "are security updates installed automatically?"
    # question mode

    # get cache
    security_pkgs = []
    try:
        cache = apt_pkg.Cache(OpNullProgress())
    except SystemError, e:
        sys.stderr.write("E: "+ _("Error: Opening the cache (%s)") % e)
        sys.exit(-1)
    depcache = apt_pkg.DepCache(cache)

    # read the pin files
    depcache.read_pinfile()
    # read the synaptic pins too
    if os.path.exists(SYNAPTIC_PINFILE):
        depcache.read_pinfile(SYNAPTIC_PINFILE)

    # init the depcache
    depcache.init()

    if depcache.broken_count > 0:
        sys.stderr.write("E: "+ _("Error: BrokenCount > 0"))
        sys.exit(-1)

    # do the upgrade (not dist-upgrade!)
    try:
        saveDistUpgrade(cache,depcache)
    except SystemError, e:
        sys.stderr.write("E: "+ _("Error: Marking the upgrade (%s)") % e)
        sys.exit(-1)

    # analyze the ugprade
    upgrades = 0
    security_updates = 0
    for pkg in cache.packages:
        # skip packages that are not marked upgraded/installed
        if not (depcache.marked_install(pkg) or depcache.marked_upgrade(pkg)):
            continue
        # check if this is really a upgrade or a false positive
        # (workaround for ubuntu #7907)
        inst_ver = pkg.current_ver
        cand_ver = depcache.get_candidate_ver(pkg)
        if cand_ver == inst_ver:
            continue

        # check for security upgrades
        upgrades = upgrades + 1	
        if isSecurityUpgrade(cand_ver):
	    security_pkgs.append(pkg.name)
            security_updates += 1
            continue

        # now check for security updates that are masked by a 
        # canidate version from another repo (-proposed or -updates)
        for ver in pkg.version_list:
            if (inst_ver and apt_pkg.version_compare(ver.ver_str, inst_ver.ver_str) <= 0):
                #print "skipping '%s' " % ver.VerStr
                continue
            if isSecurityUpgrade(ver):
	    	security_pkgs.append(pkg.name)
                security_updates += 1
                break

    # print the number of upgrades
    if options and options.show_package_names:
	if len(security_pkgs) > 0:
         	write_package_names(security_pkgs)
	else:
		print 'No updates available.'
    elif options and options.sec_pkg_count:
	if len(security_pkgs) > 0:
		print str(len(security_pkgs))  
		exit(2)
	else:
		print 0
		exit(0)
#        write_human_readable_summary(sys.stdout, security_updates)
    elif options and options.install_security_pkgs:
	if len(security_pkgs) > 0:
		r = install_security_updates(security_pkgs)
		if r != 0:
			print 'one or more updates failed to install, check apt log.'
		else:
			exit(0)
	else:
		print 'No updates available'
		exit(0)
    else:
        # print the number of regular upgrades and the number of 
        # security upgrades
	if len(security_pkgs) > 0:
		#<0|1|2|3> <checkname> count=<current_value>;<warn>;<crit>;<min>;<max> <text output of command>
		print '2 Security_updates_apt ' + 'count=' + str(security_updates) +';1;1;1;1 ' + str(security_updates) + ' pending security updates, pkg list: ' + str(security_pkgs)  
		exit(2)
	else:
		print '0 Security_updates_apt ' + 'count=' + str(security_updates) +';1;1;1;1' + ' No pending updates. '  
		exit(0)

    # return the number of upgrades (if its used as a module)
    return(upgrades,security_updates)


if __name__ == "__main__":        
    # setup a exception handler to make sure that uncaught stuff goes
    # to the notifier
    sys.excepthook = _handleException
    
    # gettext
    APP="update-notifier"
    DIR="/usr/share/locale"
    gettext.bindtextdomain(APP, DIR)
    gettext.textdomain(APP)

    # check arguments
    parser = OptionParser()
    parser.add_option("-p",
                      "--package-names",
                      action="store_true",
                      dest="show_package_names",
                      help=_("Show the security packages that are going to be installed/upgraded"))
    parser.add_option("",
                      "--security-update-count",
                      action="store_true",
                      dest="sec_pkg_count",
                      help=_("Display the count of security updates"))
    parser.add_option("",
                      "--install-security-updates",
                      action="store_true",
                      dest="install_security_pkgs",
                      help=_("install package security updates"))

    (options, args) = parser.parse_args()

    # run it
    init()
    run(options)
