#!/usr/bin/python -b
# Copyright 1999-2021 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

import argparse
import re
import sys
import codecs
from functools import reduce
import operator

from os import path as osp
if osp.isfile(osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), ".portage_not_installed")):
	sys.path.insert(0, osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), "lib"))
import portage
portage._internal_caller = True
from portage import os
from portage.glsa import match
from portage.output import green, red, nocolor, white

__program__ = "glsa-check"
__author__ = "Marius Mauch <genone@gentoo.org>"
__version__ = "1.0.1"

# option parsing
epilog = "glsa-list can contain an arbitrary number of GLSA ids," \
		" filenames containing GLSAs or the special identifiers" \
		" 'all' and 'affected'"
parser = argparse.ArgumentParser(usage=__program__ + " <option> [glsa-id | all | new | affected]",
	epilog=epilog)

modes = parser.add_argument_group("Modes")
modes.add_argument("-l", "--list", action="store_const",
		const="list", dest="mode",
		help="List a summary for the given GLSA(s) or set and whether they affect the system")
modes.add_argument("-d", "--dump", action="store_const",
		const="dump", dest="mode",
		help="Show all information about the GLSA(s) or set")
modes.add_argument("--print", action="store_const",
		const="dump", dest="mode",
		help="Alias for --dump")
modes.add_argument("-t", "--test", action="store_const",
		const="test", dest="mode",
		help="Test if this system is affected by the GLSA(s) or set and output the GLSA ID(s)")
modes.add_argument("-p", "--pretend", action="store_const",
		const="pretend", dest="mode",
		help="Show the necessary steps to remediate the system")
modes.add_argument("-f", "--fix", action="store_const",
		const="fix", dest="mode",
		help="(experimental) Attempt to remediate the system based on the instructions given in the GLSA(s) or set. This will only upgrade (when an upgrade path exists) or remove packages")
modes.add_argument("-i", "--inject", action="store_const",
		const="inject", dest="mode",
		help="Inject the given GLSA(s) into the glsa_injected file")
modes.add_argument("-m", "--mail", action="store_const",
		const="mail", dest="mode",
		help="Send a mail with the given GLSAs to the administrator")
parser.add_argument("-V", "--version", action="store_true",
		help="Show information about glsa-check")
parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
		help="Be less verbose and do not send empty mail")
parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
		help="Print more messages")
parser.add_argument("-n", "--nocolor", action="store_true",
		help="Removes color from output")
parser.add_argument("-e", "--emergelike", action="store_false", dest="least_change",
		help="Upgrade to latest version (not least-change)")
parser.add_argument("-c", "--cve", action="store_true", dest="list_cve",
		help="Show CVE IDs in listing mode")
parser.add_argument("-r", "--reverse", action="store_true", dest="reverse",
		help="List GLSAs in reverse order")

options, params = parser.parse_known_args()

if options.nocolor:
	nocolor()

if options.version:
	sys.stderr.write("\n"+ __program__ + ", version " + __version__ + "\n")
	sys.stderr.write("Author: " + __author__ + "\n")
	sys.stderr.write("This program is licensed under the GPL, version 2\n\n")
	sys.exit(0)

mode = options.mode
least_change = options.least_change
list_cve = options.list_cve
quiet = options.quiet
verbose = options.verbose

# Sanity checking
if mode is None:
	sys.stderr.write("No mode given: what should I do?\n")
	parser.print_help()
	sys.exit(1)
elif mode != "list" and not params:
	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")
	parser.print_help()
	sys.exit(1)
elif mode in ["fix", "inject"] and os.geteuid() != 0:
	# we need root privileges for write access
	sys.stderr.write("\nThis tool needs root access to "+options.mode+" this GLSA\n\n")
	sys.exit(2)
elif mode == "list" and not params:
	params.append("affected")

# delay this for speed increase
from portage.glsa import (Glsa, GlsaTypeException, GlsaFormatException,
	get_applied_glsas, get_glsa_list)

eroot = portage.settings['EROOT']
vardb = portage.db[eroot]["vartree"].dbapi
portdb = portage.db[eroot]["porttree"].dbapi

# build glsa lists
completelist = get_glsa_list(portage.settings)

checklist = get_applied_glsas(portage.settings)
todolist = [e for e in completelist if e not in checklist]

glsalist = []
if "new" in params:
	params.remove("new")
	sys.stderr.write("Warning: The 'new' glsa-list target has been removed, using 'affected'.\n")
	params.append("affected")

if "all" in params:
	glsalist = completelist
	params.remove("all")

if "affected" in params:
	for x in todolist:
		try:
			myglsa = Glsa(x, portage.settings, vardb, portdb)
		except (GlsaTypeException, GlsaFormatException) as 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"):
	# Get to the raw streams in py3k before wrapping them with an encoded writer
	# to avoid writing bytes to a text stream (stdout/stderr are text streams
	# by default in py3k)
	if hasattr(fd1, "buffer"):
		fd1 = fd1.buffer
	if hasattr(fd2, "buffer"):
		fd2 = fd2.buffer
	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")

	for myid in sorted(myglsalist, reverse=options.reverse):
		try:
			myglsa = Glsa(myid, portage.settings, vardb, portdb)
		except (GlsaTypeException, GlsaFormatException) as 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 list(myglsa.packages)[:3]:
				fd1.write(" " + pkg + " ")
			if len(myglsa.packages) > 3:
				fd1.write("... ")
		else:
			for cpv in myglsa.packages.keys():
				pkg = myglsa.packages[cpv]
				for path in pkg:
					v_installed = reduce(operator.add, [match(v, vardb) for v in path["vul_atoms"]], [])
					u_installed = reduce(operator.add, [match(u, vardb) for u in path["unaff_atoms"]], [])
				mylist = sorted(set(v_installed).difference(set(u_installed)))
				if len(mylist) > 0:
					cpv = color(" ".join(mylist))
				fd1.write(" " + cpv + " ")

		fd1.write(")")
		if list_cve:
			cve_ids = []
			for r in myglsa.references:
				m = re.search(r'(CAN|CVE)-[\d-]+', r)
				if m is not None:
					cve_ids.append(m.group(0))
			if cve_ids:
				fd1.write(" "+(",".join(cve_ids)))
		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, portage.settings, vardb, portdb)
		except (GlsaTypeException, GlsaFormatException) as 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":
			if not quiet:
				sys.stdout.write("Fixing GLSA "+myid+"\n")
			if not myglsa.isVulnerable():
				if not quiet:
					sys.stdout.write(">>> no vulnerable packages installed\n")
			else:
				if quiet:
					sys.stdout.write("Fixing GLSA "+myid+"\n")
				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" + (" --quiet" if quiet else "") + " =" + 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":
			if not quiet:
				sys.stdout.write("Checking GLSA "+myid+"\n")
			if not myglsa.isVulnerable():
				if not quiet:
					sys.stdout.write(">>> no vulnerable packages installed\n")
			else:
				if quiet:
					sys.stdout.write("Checking GLSA "+myid+"\n")
				mergedict = {}
				for (vuln, update) in myglsa.getAffectionTable(least_change=least_change):
					mergedict.setdefault(update, []).append(vuln)

				# first, extract the atoms that cannot be upgraded (where key == "")
				no_upgrades = []
				sys.stdout.write(">>> The following updates will be performed for this GLSA:\n")
				if "" in mergedict:
					no_upgrades = mergedict[""]
					del mergedict[""]

				# see if anything is left that can be upgraded
				if mergedict:
					sys.stdout.write(">>> Updates that will be performed:\n")
					for (upd, vuln) in mergedict.items():
						sys.stdout.write("     " + green(upd) + " (vulnerable: " + red(", ".join(vuln)) + ")\n")

				if no_upgrades:
					sys.stdout.write(">>> No upgrade path exists for these packages:\n")
					sys.stdout.write("     " + red(", ".join(no_upgrades)) + "\n")
			sys.stdout.write("\n")
		elif mode == "inject":
			sys.stdout.write("injecting " + myid + "\n")
			myglsa.inject()
		sys.stdout.write("\n")
	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, portage.settings, vardb, portdb)
		except (GlsaTypeException, GlsaFormatException) as 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")
			sys.exit(6)
	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":
	import socket
	from io import BytesIO
	from email.mime.text import MIMEText
	import portage.mail

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

	if "PORTAGE_ELOG_MAILURI" in portage.settings:
		myrecipient = portage.settings["PORTAGE_ELOG_MAILURI"].split()[0]
	else:
		myrecipient = "root@localhost"

	if "PORTAGE_ELOG_MAILFROM" in portage.settings:
		myfrom = portage.settings["PORTAGE_ELOG_MAILFROM"]
	else:
		myfrom = "glsa-check"

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

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

	myattachments = []
	for myid in glsalist:
		try:
			myglsa = Glsa(myid, portage.settings, vardb, portdb)
		except (GlsaTypeException, GlsaFormatException) as e:
			if verbose:
				sys.stderr.write(("invalid GLSA: %s (error message was: %s)\n" % (myid, e)))
			continue
		myfd = BytesIO()
		myglsa.dump(outstream=myfd)
		attachment = myfd.getvalue().decode("utf-8")
		myattachments.append(MIMEText(attachment, _charset="utf8"))
		myfd.close()

	if glsalist or not quiet:
		mymessage = portage.mail.create_message(myfrom, myrecipient, mysubject, summary, myattachments)
		portage.mail.send_mail(portage.settings, 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)