#!/usr/bin/python

# $Header: $
# This program is licensed under the GPL, version 2

import os
import sys
import codecs
try:
	import portage
except ImportError:
	sys.path.insert(0, "/usr/lib/portage/pym")
	import portage

try:
	from portage.output import *
except ImportError:
	from output import *

from getopt import getopt, GetoptError

__program__ = "glsa-check"
__author__ = "Marius Mauch <genone@gentoo.org>"
__version__ = open("/etc/gentoolkit-version").read().strip()

optionmap = [
["-l", "--list", "list all unapplied GLSA"],
["-d", "--dump", "--print", "show all information about the given GLSA"],
["-t", "--test", "test if this system is affected by the given GLSA"],
["-p", "--pretend", "show the necessary commands to apply this GLSA"],
["-f", "--fix", "try to auto-apply this GLSA (experimental)"],
["-i", "--inject", "inject the given GLSA into the glsa_injected file"],
["-n", "--nocolor", "disable colors (option)"],
["-e", "--emergelike", "do not use a least-change algorithm (option)"],
["-h", "--help", "show this help message"],
["-V", "--version", "some information about this tool"],
["-v", "--verbose", "print more information (option)"],
["-c", "--cve", "show CVE ids in listing mode (option)"],
["-q", "--quiet", "be less verbose and do not send empty mail (option)"],
["-m", "--mail", "send a mail with the given GLSAs to the administrator"],
]

# print a warning as this is beta code (but proven by now, so no more warning)
#sys.stderr.write("WARNING: This tool is completely new and not very tested, so it should not be\n")
#sys.stderr.write("used on production systems. It's mainly a test tool for the new GLSA release\n")
#sys.stderr.write("and distribution system, it's functionality will later be merged into emerge\n")
#sys.stderr.write("and equery.\n")
#sys.stderr.write("Please read http://www.gentoo.org/proj/en/portage/glsa-integration.xml\n")
#sys.stderr.write("before using this tool AND before reporting a bug.\n\n")

# option parsing
args = []
params = []
try:
	args, params = getopt(sys.argv[1:], "".join([o[0][1] for o in optionmap]), \
		[x[2:] for x in reduce(lambda x,y: x+y, [z[1:-1] for z in optionmap])])
	args = [a for a,b in args]
	
	for option in ["--nocolor", "-n"]:
		if option in args:
			nocolor()
			args.remove(option)
			
	verbose = False
	for option in ["--verbose", "-v"]:
		if option in args:
			verbose = True
			args.remove(option)

	list_cve = False
	for option in ["--cve", "-c"]:
		if option in args:
			list_cve = True
			args.remove(option)
	
	least_change = True
	for option in ["--emergelike", "-e"]:
		if option in args:
			least_change = False
			args.remove(option)

	quiet = False
	for option in ["--quiet", "-q"]:
		if option in args:
			quiet = True
			args.remove(option)


	# sanity checking
	if len(args) <= 0:
		sys.stderr.write("no option given: what should I do ?\n")
		mode = "HELP"
	elif len(args) > 1:
		sys.stderr.write("please use only one command per call\n")
		mode = "HELP"
	else:
		# in what mode are we ?
		args = args[0]
		for m in optionmap:
			if args in [o for o in m[:-1]]:
				mode = m[1][2:]

except GetoptError, e:
	sys.stderr.write("unknown option given: ")
	sys.stderr.write(str(e)+"\n")
	mode = "HELP"

# we need a set of glsa for most operation modes
if len(params) <= 0 and mode in ["fix", "test", "pretend", "dump", "inject", "mail"]:
	sys.stderr.write("\nno GLSA given, so we'll do nothing for now. \n")
	sys.stderr.write("If you want to run on all GLSA please tell me so \n")
	sys.stderr.write("(specify \"all\" as parameter)\n\n")
	mode = "HELP"
elif len(params) <= 0 and mode == "list":
	params.append("new")
	
# show help message
if mode == "help" or mode == "HELP":
	msg = "Syntax: glsa-check <option> [glsa-list]\n\n"
	for m in optionmap:
		msg += m[0] + "\t" + m[1] + "   \t: " + m[-1] + "\n"
		for o in m[2:-1]:
			msg += "\t" + o + "\n"
	msg += "\nglsa-list can contain an arbitrary number of GLSA ids, \n"
	msg += "filenames containing GLSAs or the special identifiers \n"
	msg += "'all', 'new' and 'affected'\n"
	if mode == "help":
		sys.stdout.write(msg)
		sys.exit(0)
	else:
		sys.stderr.write("\n" + msg)
		sys.exit(1)

# we need root privileges for write access
if mode in ["fix", "inject"] and os.geteuid() != 0:
	sys.stderr.write(__program__ + ": root access is needed for \""+mode+"\" mode\n")
	sys.exit(2)

# show version and copyright information
if mode == "version":
	sys.stderr.write("%(program)s (%(version)s)\n" % {
		"program": __program__,
		"version": __version__
	})
	sys.stderr.write("Author: %s\n" % __author__)
	sys.stderr.write("This program is licensed under the GPL, version 2\n")
	sys.exit(0)

# delay this for speed increase
from gentoolkit.glsa import *

glsaconfig = checkconfig(portage.config(clone=portage.settings))

if quiet:
    glsaconfig["EMERGE_OPTS"] += " --quiet"

vardb = portage.db["/"]["vartree"].dbapi
portdb = portage.db["/"]["porttree"].dbapi

# Check that we really have a glsa dir to work on
if not (os.path.exists(glsaconfig["GLSA_DIR"]) and os.path.isdir(glsaconfig["GLSA_DIR"])):
	sys.stderr.write(red("ERROR")+": GLSA_DIR %s doesn't exist. Please fix this.\n" % glsaconfig["GLSA_DIR"])
	sys.exit(1)

# build glsa lists
completelist = get_glsa_list(glsaconfig["GLSA_DIR"], glsaconfig)

if os.access(glsaconfig["CHECKFILE"], os.R_OK):
	checklist = [line.strip() for line in open(glsaconfig["CHECKFILE"], "r").readlines()]
else:
	checklist = []
todolist = [e for e in completelist if e not in checklist]

glsalist = []
if "new" in params:
	glsalist = todolist
	params.remove("new")
	
if "all" in params:
	glsalist = completelist
	params.remove("all")
if "affected" in params:
	# replaced completelist with todolist on request of wschlich
	for x in todolist:
		try:
			myglsa = Glsa(x, glsaconfig)
		except (GlsaTypeException, GlsaFormatException), e:
			if verbose:
				sys.stderr.write(("invalid GLSA: %s (error message was: %s)\n" % (x, e)))
			continue
		if myglsa.isVulnerable():
			glsalist.append(x)
	params.remove("affected")

# remove invalid parameters
for p in params[:]:
	if not (p in completelist or os.path.exists(p)):
		sys.stderr.write(("(removing %s from parameter list as it isn't a valid GLSA specification)\n" % p))
		params.remove(p)

glsalist.extend([g for g in params if g not in glsalist])

def summarylist(myglsalist, fd1=sys.stdout, fd2=sys.stderr, encoding="utf-8"):
	fd1 = codecs.getwriter(encoding)(fd1)
	fd2 = codecs.getwriter(encoding)(fd2)
	if not quiet:
		fd2.write(white("[A]")+" means this GLSA was marked as applied (injected),\n")
		fd2.write(green("[U]")+" means the system is not affected and\n")
		fd2.write(red("[N]")+" indicates that the system might be affected.\n\n")

	myglsalist.sort()
	for myid in myglsalist:
		try:
			myglsa = Glsa(myid, glsaconfig)
		except (GlsaTypeException, GlsaFormatException), e:
			if verbose:
				fd2.write(("invalid GLSA: %s (error message was: %s)\n" % (myid, e)))
			continue
		if myglsa.isInjected():
			status = "[A]"
			color = white
		elif myglsa.isVulnerable():
			status = "[N]"
			color = red
		else:
			status = "[U]"
			color = green

		if verbose:
			access = ("[%-8s] " % myglsa.access)
		else:
			access=""

		fd1.write(color(myglsa.nr) + " " + color(status) + " " + color(access) + myglsa.title + " (")
		if not verbose:
			for pkg in myglsa.packages.keys()[:3]:
				fd1.write(" " + pkg + " ")
			if len(myglsa.packages) > 3:
				fd1.write("... ")
		else:
			for pkg in myglsa.packages.keys():
				mylist = vardb.match(portage.dep_getkey(str(pkg)))
				if len(mylist) > 0:
					pkg = color(" ".join(mylist))
				fd1.write(" " + pkg + " ")

		fd1.write(")")
		if list_cve:
			fd1.write(" "+(",".join([r[:13] for r in myglsa.references if r[:4] in ["CAN-", "CVE-"]])))
		fd1.write("\n")
	return 0

if mode == "list":
	sys.exit(summarylist(glsalist))

# dump, fix, inject and fix are nearly the same code, only the glsa method call differs
if mode in ["dump", "fix", "inject", "pretend"]:
	for myid in glsalist:
		try:
			myglsa = Glsa(myid, glsaconfig)
		except (GlsaTypeException, GlsaFormatException), e:
			if verbose:
				sys.stderr.write(("invalid GLSA: %s (error message was: %s)\n" % (myid, e)))
			continue
		if mode == "dump":
			myglsa.dump()
		elif mode == "fix":
			sys.stdout.write("Fixing GLSA "+myid+"\n")
			if not myglsa.isVulnerable():
				sys.stdout.write(">>> no vulnerable packages installed\n")
			else:
				mergelist = myglsa.getMergeList(least_change=least_change)
				if mergelist == []:
					sys.stdout.write(">>> cannot fix GLSA, no unaffected packages available\n")
					sys.exit(2)
				for pkg in mergelist:
					sys.stdout.write(">>> merging "+pkg+"\n")
					# using emerge for the actual merging as it contains the dependency
					# code and we want to be consistent in behaviour. Also this functionality
					# will be integrated in emerge later, so it shouldn't hurt much.
					emergecmd = "emerge --oneshot " + glsaconfig["EMERGE_OPTS"] + " =" + pkg
					if verbose:
						sys.stderr.write(emergecmd+"\n")
					exitcode = os.system(emergecmd)
					# system() returns the exitcode in the high byte of a 16bit integer
					if exitcode >= 1<<8:
						exitcode >>= 8
					if exitcode:
						sys.exit(exitcode)
			if len(mergelist):
				sys.stdout.write("\n")
		elif mode == "pretend":
			sys.stdout.write("Checking GLSA "+myid+"\n")
			if not myglsa.isVulnerable():
				sys.stdout.write(">>> no vulnerable packages installed\n")
			else:
				mergedict = {}
				for (vuln, update) in myglsa.getAffectionTable(least_change=least_change):
					mergedict.setdefault(update, []).append(vuln)
				
				sys.stdout.write(">>> The following updates will be performed for this GLSA:\n")
				for pkg in mergedict:
					if pkg != "":
						sys.stdout.write("     " + pkg + " (vulnerable: " + ", ".join(mergedict[pkg]) + ")\n")
				if "" in mergedict:
					sys.stdout.write("\n>>> For the following packages, no upgrade path exists:\n")
					sys.stdout.write("     " + ", ".join(mergedict[""]))
		elif mode == "inject":
			sys.stdout.write("injecting " + myid + "\n")
			myglsa.inject()
	sys.exit(0)

# test is a bit different as Glsa.test() produces no output
if mode == "test":
	outputlist = []
	for myid in glsalist:
		try:
			myglsa = Glsa(myid, glsaconfig)
		except (GlsaTypeException, GlsaFormatException), e:
			if verbose:
				sys.stderr.write(("invalid GLSA: %s (error message was: %s)\n" % (myid, e)))
			continue
		if myglsa.isVulnerable():
			outputlist.append(str(myglsa.nr))
	if len(outputlist) > 0:
		sys.stderr.write("This system is affected by the following GLSAs:\n")
		if verbose:
			summarylist(outputlist)
		else:
			sys.stdout.write("\n".join(outputlist)+"\n")
	else:
		sys.stderr.write("This system is not affected by any of the listed GLSAs\n")
	sys.exit(0)

# mail mode as requested by solar
if mode == "mail":
	try:
		import portage.mail as portage_mail
	except ImportError:
		import portage_mail
		
	import socket
	from StringIO import StringIO
	try:
		from email.mime.text import MIMEText
	except ImportError:
		from email.MIMEText import MIMEText

	# color doesn't make any sense for mail
	nocolor()

	if "PORTAGE_ELOG_MAILURI" in glsaconfig:
		myrecipient = glsaconfig["PORTAGE_ELOG_MAILURI"].split()[0]
	else:
		myrecipient = "root@localhost"
	
	if "PORTAGE_ELOG_MAILFROM" in glsaconfig:
		myfrom = glsaconfig["PORTAGE_ELOG_MAILFROM"]
	else:
		myfrom = "glsa-check"

	mysubject = "[glsa-check] Summary for %s" % socket.getfqdn()

	# need a file object for summarylist()
	myfd = StringIO()
	myfd.write("GLSA Summary report for host %s\n" % socket.getfqdn())
	myfd.write("(Command was: %s)\n\n" % " ".join(sys.argv))
	summarylist(glsalist, fd1=myfd, fd2=myfd)
	summary = str(myfd.getvalue())
	myfd.close()

	myattachments = []
	for myid in glsalist:
		try:
			myglsa = Glsa(myid, glsaconfig)
		except (GlsaTypeException, GlsaFormatException), e:
			if verbose:
				sys.stderr.write(("invalid GLSA: %s (error message was: %s)\n" % (myid, e)))
			continue
		myfd = StringIO()
		myglsa.dump(outstream=myfd)
		myattachments.append(MIMEText(str(myfd.getvalue()), _charset="utf8"))
		myfd.close()

        if glsalist or not quiet:
		mymessage = portage_mail.create_message(myfrom, myrecipient, mysubject, summary, myattachments)
		portage_mail.send_mail(glsaconfig, mymessage)
		
	sys.exit(0)
	
# something wrong here, all valid paths are covered with sys.exit()
sys.stderr.write("nothing more to do\n")
sys.exit(2)