#!/usr/bin/python3
#
#   40-updates - create the list of packages for update with caching
#   Copyright (c) 2015 Igor Pecovnik
#
#   Author:  Igor Pecovnik igor.pecovnik@gmail.com
#   Based upon prior work by Nick Charlton, Dustin Kirkland and Michael Vogt.
#
#   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, write to the Free Software Foundation, Inc.,
#   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import sys
import subprocess
import apt_pkg
import time
import os
import os.path as path

def refreshdata():
		# update procedure
		DISTRO = subprocess.Popen(["lsb_release", "-c", "-s"],
								  stdout=subprocess.PIPE).communicate()[0].strip()

		class OpNullProgress(object):
			'''apt progress handler which supresses any output.'''
			def update(self):
				pass
			def done(self):
				pass

		def is_security_upgrade(pkg):
			'''
			Checks to see if a package comes from a DISTRO-security source.
			'''
			security_package_sources = [("Ubuntu", "%s-security" % DISTRO),
									   ("Debian", "%s-security" % DISTRO)]

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

		# init apt and config
		apt_pkg.init()

		# open the apt cache
		try:
			cache = apt_pkg.Cache(OpNullProgress())
		except SystemError as e:
			sys.stderr.write("Error: Opening the cache (%s)" % e)
			sys.exit(0)

		# setup a DepCache instance to interact with the repo
		depcache = apt_pkg.DepCache(cache)

		# take into account apt policies
		depcache.read_pinfile()

		# initialise it
		depcache.init()

		# give up if packages are broken
		if depcache.broken_count > 0:
			sys.stderr.write("Error: Broken packages exist.")
			sys.exit(0)

		# mark possible packages
		try:
			# run distro-upgrade
			depcache.upgrade(True)
			# reset if packages get marked as deleted -> we don't want to break anything
			if depcache.del_count > 0:
				depcache.init()

			# then a standard upgrade
			depcache.upgrade()
		except SystemError as e:
			sys.stderr.write("Error: Couldn't mark the upgrade (%s)" % e)
			sys.exit(0)

		# run around the packages
		upgrades = 0
		security_upgrades = 0
		for pkg in cache.packages:
			candidate = depcache.get_candidate_ver(pkg)
			current = pkg.current_ver

			# skip packages not marked as upgraded/installed
			if not (depcache.marked_install(pkg) or depcache.marked_upgrade(pkg)):
				continue

			# increment the upgrade counter
			upgrades += 1

			# keep another count for security upgrades
			if is_security_upgrade(candidate):
				security_upgrades += 1

			# double check for security upgrades masked by another package
			for version in pkg.version_list:
				if (current and apt_pkg.version_compare(version.ver_str, current.ver_str) <= 0):
					continue
				if is_security_upgrade(version):
					security_upgrades += 1
					break

		# to properly update time stamp we need to remove file first
		if os.path.isfile(myfile):
			os.remove(myfile)

		f = open(myfile,'w')
		f.write('%d' % (upgrades))
		f.close() # you can omit in most cases as the destructor will call it else:
		os.utime(myfile, None)
		return


# if the file is older than 2 days
myfile = "/var/cache/apt/archives/updates.number"
now = time.time()
twodays_ago = now - 60*60*24*2 # Number of seconds in two days
lastupgrade = path.getmtime('/var/lib/dpkg/status')

if path.isfile(myfile):
	fileCreation = path.getmtime(myfile)
	if (fileCreation < twodays_ago or lastupgrade > fileCreation):
		refreshdata()
else:
		refreshdata()

# display if there are some upgrades
file = open(myfile, 'r+')
updates = int(file.read())
if updates > 0:
	print("\n[ \033[92m%d updates to install\033[0m: apt-get upgrade ]\n" % (updates))