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

from __future__ import division, print_function

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

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__))), "pym"))
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 ConfigProtect, ensure_dirs, shlex_split, varexpand, _xattr
xattr = _xattr.xattr
from portage.dbapi.vartree import dblink, tar_contents
from portage.checksum import perform_md5
from portage._sets import load_default_config, SETPREFIX
from portage.process import find_binary
from portage.util.compression_probe import _compressors

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

	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
			contents = dblnk.getcontents()
			protect = None
			if not include_config:
				confprot = ConfigProtect(eroot,
					shlex_split(settings.get("CONFIG_PROTECT", "")),
					shlex_split(settings.get("CONFIG_PROTECT_MASK", "")),
					case_insensitive=("case-insensitive-fs"
					in settings.features))
				def protect(filename):
					if not confprot.isprotected(filename):
						return False
					if include_unmodified_config:
						file_data = contents[filename]
						if file_data[0] == "obj":
							orig_md5 = file_data[2].lower()
							cur_md5 = perform_md5(filename, calc_prelink=1)
							if orig_md5 == cur_md5:
								return False
					excluded_config_files.append(filename)
					return True
			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(os.getpid()))
			ensure_dirs(os.path.dirname(binpkg_tmpfile))
			binpkg_compression = settings.get("BINPKG_COMPRESS", "bzip2")
			try:
				compression = _compressors[binpkg_compression]
			except KeyError as e:
				eout.eerror("Invalid or unsupported compression method: %s" % e.args[0])
				return 1
			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 = [varexpand(x, mydict=settings) for x in shlex_split(compression["compress"])]
			# 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)
				# The tarfile module will write pax headers holding the
				# xattrs only if PAX_FORMAT is specified here.
				with tarfile.open(mode="w|",format=tarfile.PAX_FORMAT if xattrs else tarfile.DEFAULT_FORMAT, fileobj=proc.stdin) as tar:
					tar_contents(contents, root, tar, protect=protect, xattrs=xattrs)
				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)
	sys.exit(retval)