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

import argparse
import errno
import math
import signal
import subprocess
import sys

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 import xpak
from portage.dbapi.dep_expand import dep_expand
from portage.dep import Atom, use_reduce
from portage.exception import (AmbiguousPackageName, InvalidAtom, InvalidData,
	InvalidDependString, PackageSetNotFound, PermissionDenied)
from portage.util import ensure_dirs, shlex_split, varexpand, _xattr
xattr = _xattr.xattr
from portage._sets import load_default_config, SETPREFIX
from portage.process import find_binary
from portage.util.compression_probe import _compressors
from portage.util._eventloop.global_event_loop import global_event_loop


def quickpkg_atom(options, infos, arg, eout):
	settings = portage.settings
	eroot = portage.settings['EROOT']
	trees = portage.db[eroot]
	vartree = trees["vartree"]
	vardb = vartree.dbapi
	bintree = trees["bintree"]

	include_config = options.include_config == "y"
	include_unmodified_config = options.include_unmodified_config == "y"
	fix_metadata_keys = ["PF", "CATEGORY"]

	try:
		atom = dep_expand(arg, mydb=vardb, settings=vartree.settings)
	except AmbiguousPackageName as e:
		# Multiple matches thrown from cpv_expand
		eout.eerror("Please use a more specific atom: %s" % \
			" ".join(e.args[0]))
		del e
		infos["missing"].append(arg)
		return 1
	except (InvalidAtom, InvalidData):
		eout.eerror("Invalid atom: %s" % (arg,))
		infos["missing"].append(arg)
		return 1
	if atom[:1] == '=' and arg[:1] != '=':
		# dep_expand() allows missing '=' but it's really invalid
		eout.eerror("Invalid atom: %s" % (arg,))
		infos["missing"].append(arg)
		return 1

	matches = vardb.match(atom)
	pkgs_for_arg = 0
	retval = 0
	for cpv in matches:
		excluded_config_files = []
		dblnk = vardb._dblink(cpv)
		have_lock = False

		if "__PORTAGE_INHERIT_VARDB_LOCK" not in settings:
			try:
				dblnk.lockdb()
				have_lock = True
			except PermissionDenied:
				pass

		try:
			if not dblnk.exists():
				# unmerged by a concurrent process
				continue
			iuse, use, restrict = vardb.aux_get(cpv,
				["IUSE","USE","RESTRICT"])
			iuse = [ x.lstrip("+-") for x in iuse.split() ]
			use = use.split()
			try:
				restrict = use_reduce(restrict, uselist=use, flat=True)
			except InvalidDependString as e:
				eout.eerror("Invalid RESTRICT metadata " + \
					"for '%s': %s; skipping" % (cpv, str(e)))
				del e
				continue
			if "bindist" in iuse and "bindist" not in use:
				eout.ewarn("%s: package was emerged with USE=-bindist!" % cpv)
				eout.ewarn("%s: it might not be legal to redistribute this." % cpv)
			elif "bindist" in restrict:
				eout.ewarn("%s: package has RESTRICT=bindist!" % cpv)
				eout.ewarn("%s: it might not be legal to redistribute this." % cpv)
			eout.ebegin("Building package for %s" % cpv)
			pkgs_for_arg += 1
			existing_metadata = dict(zip(fix_metadata_keys,
				vardb.aux_get(cpv, fix_metadata_keys)))
			category, pf = portage.catsplit(cpv)
			required_metadata = {}
			required_metadata["CATEGORY"] = category
			required_metadata["PF"] = pf
			update_metadata = {}
			for k, v in required_metadata.items():
				if v != existing_metadata[k]:
					update_metadata[k] = v
			if update_metadata:
				vardb.aux_update(cpv, update_metadata)
			xpdata = xpak.xpak(dblnk.dbdir)
			binpkg_tmpfile = os.path.join(bintree.pkgdir,
				cpv + ".tbz2." + str(portage.getpid()))
			ensure_dirs(os.path.dirname(binpkg_tmpfile))
			binpkg_compression = settings.get("BINPKG_COMPRESS", "bzip2")
			try:
				compression = _compressors[binpkg_compression]
			except KeyError as e:
				if binpkg_compression:
					eout.eerror("Invalid or unsupported compression method: %s" % e.args[0])
					return 1
				# Empty BINPKG_COMPRESS disables compression.
				binpkg_compression = 'none'
				compression = {
					'compress': 'cat',
					'package': 'sys-apps/coreutils',
				}
			try:
				compression_binary = shlex_split(varexpand(compression["compress"], mydict=settings))[0]
			except IndexError as e:
				eout.eerror("Invalid or unsupported compression method: %s" % e.args[0])
				return 1
			if find_binary(compression_binary) is None:
				missing_package = compression["package"]
				eout.eerror("File compression unsupported %s. Missing package: %s" % (binpkg_compression, missing_package))
				return 1
			cmd = shlex_split(varexpand(compression["compress"], mydict=settings))
			# Filter empty elements that make Popen fail
			cmd = [x for x in cmd if x != ""]
			with open(binpkg_tmpfile, "wb") as fobj:
				proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=fobj)
				excluded_config_files = dblnk.quickpkg(proc.stdin,
					include_config=include_config,
					include_unmodified_config=include_unmodified_config)
				proc.stdin.close()
				if proc.wait() != os.EX_OK:
					eout.eend(1)
					eout.eerror("Compressor failed for package %s" % cpv)
					retval |= 1
					try:
						os.unlink(binpkg_tmpfile)
					except OSError as e:
						if e.errno not in (errno.ENOENT, errno.ESTALE):
							raise
					continue
			xpak.tbz2(binpkg_tmpfile).recompose_mem(xpdata)
		finally:
			if have_lock:
				dblnk.unlockdb()
		pkg_info = bintree.inject(cpv, filename=binpkg_tmpfile)
		# The pkg_info value ensures that the following getname call
		# returns the correct path when FEATURES=binpkg-multi-instance
		# is enabled, but fallback to cpv in case the inject call
		# returned None due to some kind of failure.
		binpkg_path = bintree.getname(pkg_info or cpv)
		try:
			s = os.stat(binpkg_path)
		except OSError:
			s = None

		if s is None or pkg_info is None:
			# Sanity check, shouldn't happen normally.
			eout.eend(1)
			eout.eerror("Failed to create package: '%s'" % binpkg_path)
			retval |= 1
		else:
			eout.eend(0)
			infos["successes"].append((cpv, s.st_size))
			infos["config_files_excluded"] += len(excluded_config_files)
			for filename in excluded_config_files:
				eout.ewarn("Excluded config: '%s'" % filename)
	if not pkgs_for_arg:
		eout.eerror("Could not find anything " + \
			"to match '%s'; skipping" % arg)
		infos["missing"].append(arg)
		retval |= 1
	return retval

def quickpkg_set(options, infos, arg, eout):
	eroot = portage.settings['EROOT']
	trees = portage.db[eroot]
	vartree = trees["vartree"]

	settings = vartree.settings
	settings._init_dirs()
	setconfig = load_default_config(settings, trees)
	sets = setconfig.getSets()

	set = arg[1:]
	if not set in sets:
		eout.eerror("Package set not found: '%s'; skipping" % (arg,))
		infos["missing"].append(arg)
		return 1

	try:
		atoms = setconfig.getSetAtoms(set)
	except PackageSetNotFound as e:
		eout.eerror("Failed to process package set '%s' because " % set +
			"it contains the non-existent package set '%s'; skipping" % e)
		infos["missing"].append(arg)
		return 1
	retval = os.EX_OK
	for atom in atoms:
		retval |= quickpkg_atom(options, infos, atom, eout)
	return retval


def quickpkg_extended_atom(options, infos, atom, eout):
	eroot = portage.settings['EROOT']
	trees = portage.db[eroot]
	vartree = trees["vartree"]
	vardb = vartree.dbapi

	require_metadata = atom.slot or atom.repo
	atoms = []
	for cpv in vardb.cpv_all():
		cpv_atom = Atom("=%s" % cpv)

		if atom == "*/*":
			atoms.append(cpv_atom)
			continue

		if not portage.match_from_list(atom, [cpv]):
			continue

		if require_metadata:
			try:
				cpv = vardb._pkg_str(cpv, atom.repo)
			except (KeyError, InvalidData):
				continue
			if not portage.match_from_list(atom, [cpv]):
				continue

		atoms.append(cpv_atom)

	for atom in atoms:
		quickpkg_atom(options, infos, atom, eout)


def quickpkg_main(options, args, eout):
	eroot = portage.settings['EROOT']
	trees = portage.db[eroot]
	bintree = trees["bintree"]

	try:
		ensure_dirs(bintree.pkgdir)
	except portage.exception.PortageException:
		pass
	if not os.access(bintree.pkgdir, os.W_OK):
		eout.eerror("No write access to '%s'" % bintree.pkgdir)
		return errno.EACCES

	if 'xattr' in portage.settings.features and not _xattr.XATTRS_WORKS:
		eout.eerror("No xattr support library was found, "
			"so xattrs will not be preserved!")
		portage.settings.unlock()
		portage.settings.features.remove('xattr')
		portage.settings.lock()

	infos = {}
	infos["successes"] = []
	infos["missing"] = []
	infos["config_files_excluded"] = 0
	for arg in args:
		if arg[0] == SETPREFIX:
			quickpkg_set(options, infos, arg, eout)
			continue
		try:
			atom = Atom(arg, allow_wildcard=True, allow_repo=True)
		except (InvalidAtom, InvalidData):
			# maybe it's valid but missing category (requires dep_expand)
			quickpkg_atom(options, infos, arg, eout)
		else:
			if atom.extended_syntax:
				quickpkg_extended_atom(options, infos, atom, eout)
			else:
				quickpkg_atom(options, infos, atom, eout)

	if not infos["successes"]:
		eout.eerror("No packages found")
		return 1
	print()
	eout.einfo("Packages now in '%s':" % bintree.pkgdir)
	units = {10:'K', 20:'M', 30:'G', 40:'T',
		50:'P', 60:'E', 70:'Z', 80:'Y'}
	for cpv, size in infos["successes"]:
		if not size:
			# avoid OverflowError in math.log()
			size_str = "0"
		else:
			power_of_2 = math.log(size, 2)
			power_of_2 = 10*(power_of_2//10)
			unit = units.get(power_of_2)
			if unit:
				size = float(size)/(2**power_of_2)
				size_str = "%.1f" % size
				if len(size_str) > 4:
					# emulate `du -h`, don't show too many sig figs
					size_str = str(int(size))
				size_str += unit
			else:
				size_str = str(size)
		eout.einfo("%s: %s" % (cpv, size_str))
	if infos["config_files_excluded"]:
		print()
		eout.ewarn("Excluded config files: %d" % infos["config_files_excluded"])
		eout.ewarn("See --help if you would like to include config files.")
	if infos["missing"]:
		print()
		eout.ewarn("The following packages could not be found:")
		eout.ewarn(" ".join(infos["missing"]))
		return 2
	return os.EX_OK

if __name__ == "__main__":
	usage = "quickpkg [options] <list of package atoms or package sets>"
	parser = argparse.ArgumentParser(usage=usage)
	parser.add_argument("--umask",
		default="0077",
		help="umask used during package creation (default is 0077)")
	parser.add_argument("--ignore-default-opts",
		action="store_true",
		help="do not use the QUICKPKG_DEFAULT_OPTS environment variable")
	parser.add_argument("--include-config",
		choices=["y","n"],
		default="n",
		metavar="<y|n>",
		help="include all files protected by CONFIG_PROTECT (as a security precaution, default is 'n')")
	parser.add_argument("--include-unmodified-config",
		choices=["y","n"],
		default="n",
		metavar="<y|n>",
		help="include files protected by CONFIG_PROTECT that have not been modified since installation (as a security precaution, default is 'n')")
	options, args = parser.parse_known_args(sys.argv[1:])
	if not options.ignore_default_opts:
		default_opts = shlex_split(
			portage.settings.get("QUICKPKG_DEFAULT_OPTS", ""))
		options, args = parser.parse_known_args(default_opts + sys.argv[1:])
	if not args:
		parser.error("no packages atoms given")
	try:
		umask = int(options.umask, 8)
	except ValueError:
		parser.error("invalid umask: %s" % options.umask)
	# We need to ensure a sane umask for the packages that will be created.
	old_umask = os.umask(umask)
	eout = portage.output.EOutput()
	def sigwinch_handler(signum, frame):
		lines, eout.term_columns =  portage.output.get_term_size()
	signal.signal(signal.SIGWINCH, sigwinch_handler)
	try:
		retval = quickpkg_main(options, args, eout)
	finally:
		os.umask(old_umask)
		signal.signal(signal.SIGWINCH, signal.SIG_DFL)
		global_event_loop().close()
	sys.exit(retval)