#!/usr/bin/python2.6

#

# CDDL HEADER START

#

# The contents of this file are subject to the terms of the

# Common Development and Distribution License (the "License").

# You may not use this file except in compliance with the License.

#

# You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE

# or http://www.opensolaris.org/os/licensing.

# See the License for the specific language governing permissions

# and limitations under the License.

#

# When distributing Covered Code, include this CDDL HEADER in each

# file and include the License file at usr/src/OPENSOLARIS.LICENSE.

# If applicable, add the following below this CDDL HEADER, with the

# fields enclosed by brackets "[]" replaced with your own identifying

# information: Portions Copyright [yyyy] [name of copyright owner]

#

# CDDL HEADER END

#

# Copyright (c) 2011, Oracle and/or its affiliates. All rights reserved.

#

try:

        import calendar

        import collections

        import getopt

        import gettext

        import itertools

        import locale

        import os

        import shutil

        import sys

        import tempfile

        import traceback

        import urllib

        import pkg.actions as actions

        import pkg.fmri

        import pkg.client.api_errors as apx

        import pkg.client.progress as progress

        import pkg.client.publisher as publisher

        import pkg.client.transport.transport as transport

        import pkg.manifest as manifest

        import pkg.misc as misc

        import pkg.publish.transaction as trans

        from pkg.misc import PipeError, emsg, msg

except KeyboardInterrupt:

        import sys

        sys.exit(1)

BUILD_RELEASE  = "5.11"  # Should be an option to this program some day?

catalog_dict   = {}    # hash table of catalogs by source uri

fmri_cache     = {}

manifest_cache = {}

null_manifest  = manifest.Manifest()

tmpdir         = None

dry_run        = False

xport          = None

dest_xport     = None

pubs           = set()

target_pub     = None

def cleanup():

        """To be called at program finish."""

        if tmpdir:

                shutil.rmtree(tmpdir, ignore_errors=True)

        if dry_run:

                return

        # Attempt to kick off a refresh of target repository for each

        # publisher before exiting.

        for pfx in pubs:

                target_pub.prefix = pfx

                try:

                        dest_xport.publish_refresh_packages(target_pub)

                except apx.TransportError:

                        # If this fails, ignore it as this was a last

                        # ditch attempt anyway.

                        break

def usage(errmsg="", exitcode=2):

        """Emit a usage message and optionally prefix it with a more specific

        error message.  Causes program to exit."""

        if errmsg:

                emsg("pkgmerge: %s" % errmsg)

        msg(_("""\

Usage:

        pkgmerge [-n] -d dest_repo -s variant=value[,...],src_repo ...

            [pkg_fmri_pattern ...]

Options:

        -d dest_repo

                The filesystem path or URI of the target repository to publish

                the merged packages to.  The target repository must already

                exist; new repositories can be created using pkgrepo(1).

        -n

                Perform a trial run with no changes made to the target

                repository.

        -s variant=value,src_repo

                The variant name and value to use for packages from this source,

                followed by the filesystem path or URI of the source repository

                or package archive to retrieve packages from.  Multiple variants

                may be specified separated by commas.  The same variants must

                be named for all sources.  This option may be specified multiple

                times.

        --help or -?

                Displays a usage message.

Environment:

        TMPDIR, TEMP, TMP

                The absolute path of the directory where temporary data should

                be stored during program execution.

"""))

        sys.exit(exitcode)

def error(text, exitcode=1):

        """Emit an error message prefixed by the command name """

        emsg("pkgmerge: %s" % text)

        if exitcode != None:

                sys.exit(exitcode)

def get_tracker(quiet=False):

        if quiet:

                progresstracker = progress.QuietProgressTracker()

        else:

                try:

                        progresstracker = \

                            progress.FancyUNIXProgressTracker()

                except progress.ProgressTrackerException:

                        progresstracker = progress.CommandLineProgressTracker()

        return progresstracker

def load_catalog(repouri, pub):

        """Load catalog from specified uri"""

        # Pull catalog only from this host

        pub.refresh(full_refresh=True, immediate=True)

        catalog_dict[repouri.uri] = dict(

            (name, [

                entry

                for entry in pub.catalog.fmris_by_version(name)

            ])

            for name in pub.catalog.names()

        )

        # Discard catalog.

        pub.remove_meta_root()

        # XXX At the moment, the only way to force the publisher object to

def get_all_pkg_names(repouri):

        return catalog_dict[repouri.uri].keys()

def get_manifest(repouri, fmri):

        """Fetch the manifest for package-fmri 'fmri' from the source

        in 'repouri'... return as Manifest object."""

        # support null manifests to keep lists ordered for merge

        if not fmri:

                return null_manifest

        mfst_str = xport.get_manifest(fmri, pub=repouri, content_only=True)

        m = manifest.Manifest(fmri)

        m.set_content(content=mfst_str)

        return m

def main_func():

        global dry_run, tmpdir, xport, dest_xport, target_pub

        dest_repo     = None

        source_list   = []

        variant_list  = []

        try:

                opts, pargs = getopt.getopt(sys.argv[1:], "d:ns:?", ["help"])

                for opt, arg in opts:

                        if opt == "-d":

                                dest_repo = misc.parse_uri(arg)

                        elif opt == "-n":

                                dry_run = True

                        elif opt == "-s":

                                s = arg.split(",")

                                if len(s) < 2:

                                        usage("-s option must specify "

                                            "variant=value,repo_uri")

                                # All but last part should be variant.

                                src_vars = {}

                                for v in s[:-1]:

                                        try:

                                                vname, vval = v.split("=")

                                        except ValueError:

                                                usage("-s option must specify "

                                                    "variant=value,repo_uri")

                                        if not vname.startswith("variant."):

                                                vname = "variant.%s" % vname

                                        src_vars[vname] = vval

                                variant_list.append(src_vars)

                                source_list.append(publisher.RepositoryURI(

                                    misc.parse_uri(s[-1])))

                        if opt in ("--help", "-?"):

                                usage(exitcode=0)

        except getopt.GetoptError, e:

                usage(_("illegal option -- %s") % e.opt)

        if not source_list:

                usage(_("At least one variant name, value, and package source "

                   "must be provided using -s."))

        if not dest_repo:

                usage(_("A destination package repository must be provided "

                    "using -d."))

        # Determine the unique set of variants across all sources.

        variants = set()

        vcombos = collections.defaultdict(set)

        for src_vars in variant_list:

                for v, vval in src_vars.iteritems():

                        variants.add(v)

                        vcombos[v].add((v, vval))

        # merge_fmris() expects this to be a list.

        variants = list(variants)

        # Require that the user specified the same variants for all sources.

        for i, src_vars in enumerate(variant_list):

                missing = set(v for v in variants if v not in variant_list[i])

                if missing:

                        missing = ", ".join(missing)

                        source = source_list[i]

                        usage(_("Source %(source)s missing values for "

                            "variants: %(missing)s") % locals())

        # Require that each unique variant combination has a source.

        for combo in itertools.product(*vcombos.values()):

                found = False

                for i, src in enumerate(source_list):

                        for vname, vval in combo:

                                if variant_list[i].get(vname,

                                    None) != vval:

                                        found = False

                                        break

                        else:

                                found = True

                                break

                if not found:

                        combo = " ".join(

                            "%s=%s" % (vname, vval)

                            for vname, vval in combo

                        )

                        usage(_("No source was specified for variant "

                            "combination %(combo)s.") % locals())

        # initialize transport

        # we use a single endpoint for now, since the transport code

        # uses publisher as a unique key... so we just flop the repo

        # list as needed to access the different catalogs/manifests/files.

        temp_root = misc.config_temp_root()

        tmpdir = tempfile.mkdtemp(dir=temp_root, prefix="pkgmerge")

        xport, xport_cfg = transport.setup_transport()

        xport_cfg.incoming_root = tmpdir

        pub = transport.setup_publisher(source_list,

            "pkgmerge", xport, xport_cfg, remote_prefix=True)

        cat_dir = tempfile.mkdtemp(dir=tmpdir)

        pub.meta_root = cat_dir

        pub.transport = xport

        # Use separate transport for destination repository in case source

        # and destination have identical publisher configuration.

        dest_xport, dest_xport_cfg = transport.setup_transport()

        dest_xport_cfg.incoming_root = tmpdir

        # retrieve catalogs for all specified repositories

        for s in source_list:

                load_catalog(s, pub)

        # determine the list of packages we'll be processing

        if not pargs:

                # use the latest versions and merge everything

                fmri_arguments = list(set(

                    name

                    for s in source_list

                    for name in get_all_pkg_names(s)

                ))

                exclude_args = []

        else:

                fmri_arguments = [

                    f

                    for f in pargs

                    if not f.startswith("!")

                ]

                exclude_args = [

                    f[1:]

                    for f in pargs

                    if f.startswith("!")

                ]

        # build fmris to be merged

        masterlist = [

            build_merge_list(fmri_arguments, exclude_args, catalog_dict[s.uri])

            for s in source_list

        ]

        # check for unmatched patterns

        in_none = reduce(lambda x, y: x & y, (set(u) for d, u in masterlist))

        if in_none:

                error(_("The following pattern(s) did not match any packages "

                    "in any of the specified repositories:\n%s") % "\n".join(

                    in_none))

        # generate set of all package names to be processed, and dict of lists

        # indexed by order in source_list; if that repo has no fmri for this

        # pkg then use None.

        allpkgs = set(name for d, u in masterlist for name in d)

        processdict = {}

        for p in allpkgs:

                for d, u in masterlist:

                        processdict.setdefault(p, []).append(d.setdefault(p,

                            None))

        # check to make sure all fmris are at same version modulo timestamp

        for entry in processdict:

                if len(set([

                        str(a).rsplit(":")[0]

                        for a in processdict[entry]

                        if a is not None

                    ])) > 1:

                        error(_("fmris matching the following patterns do not "

                            "have matching versions across all repositories: "

                            "%s") % processdict[entry])

        # we're ready to merge

        if not dry_run:

                target_pub = transport.setup_publisher(dest_repo,

                    pub.prefix, dest_xport, dest_xport_cfg, remote_prefix=True)

        def get_basename(pfmri):

                open_time = pfmri.get_timestamp()

                return "%d_%s" % \

                    (calendar.timegm(open_time.utctimetuple()),

                    urllib.quote(str(pfmri), ""))

        tracker = get_tracker()

        tracker.republish_set_goal(len(processdict), 0, 0)

        for entry in processdict:

                man, retrievals = merge_fmris(source_list,

                    processdict[entry], variant_list, variants)

                # Determine total bytes to retrieve for this package; this must

                # be done using the retrievals dict since they are coalesced by

                # hash.

                getbytes = sum(

                    misc.get_pkg_otw_size(a)

                    for i, uri in enumerate(source_list)

                    for a in retrievals[i]

                )

                # Determine total bytes to send for this package; this must be

                # done using the manifest since retrievals are coalesced based

                # on hash, but sends are not.

                sendbytes = sum(

                    int(a.attrs.get("pkg.size", 0))

                    for a in man.gen_actions()

                )

                f = man.fmri

                tracker.republish_start_pkg(f.pkg_name, getbytes=getbytes,

                    sendbytes=sendbytes)

                if dry_run:

                        # Dry-run; attempt a merge of everything but don't

                        # write any data or publish packages.

                        tracker.republish_end_pkg()

                        continue

                pkgdir = tempfile.mkdtemp(dir=tmpdir)

                # Retrieve package data from each package source.

                for i, uri in enumerate(source_list):

                        mfile = xport.multi_file_ni(pub, pkgdir,

                            decompress=True, progtrack=tracker)

                        for a in retrievals[i]:

                                mfile.add_action(a)

                        mfile.wait_files()

                # Publish merged package.

                try:

                        trans_id = get_basename(f)

                        pkg_name = f.get_fmri()

                        target_pub.prefix = f.publisher

                        pubs.add(f.publisher)

                        t = trans.Transaction(dest_repo, pkg_name=pkg_name,

                            trans_id=trans_id, xport=dest_xport, pub=target_pub,

                            progtrack=tracker)

                        # Remove any previous failed attempt to

                        # to republish this package.

                        try:

                                t.close(abandon=True)

                        except:

                                # It might not exist already.

                                pass

                        t.open()

                        for a in man.gen_actions():

                                if (a.name == "set" and

                                    a.attrs["name"] == "pkg.fmri"):

                                        # To be consistent with the

                                        # server, the fmri can't be

                                        # added to the manifest.

                                        continue

                                if hasattr(a, "hash"):

                                        fname = os.path.join(pkgdir,

                                            a.hash)

                                        a.data = lambda: open(fname, "rb")

                                t.add(a)

                        # Always defer catalog update.

                        t.close(add_to_catalog=False)

                except trans.TransactionError, e:

                        error(str(e))

                # Done with this package.

                tracker.republish_end_pkg()

                # Dump retrieved package data after each republication.

                shutil.rmtree(pkgdir)

        tracker.republish_done()

        tracker.reset()

        return 0

def merge_fmris(source_list, fmri_list, variant_list, variants):

        """Merge a list of manifests representing multiple variants,

        returning the merged manifest and a list of lists of actions to

        retrieve from each source"""

        # Merge each variant one at a time.

        merged = {}

        # where to find files...

        hash_source = {}

        for i, variant in enumerate(variants):

                # Build the unique list of remaining variant combinations to

                # use for merging this variant.

                combos = set(

                    tuple(

                        (vname, src_vars[vname])

                        for vname in variants[i + 1:]

                    )

                    for src_vars in variant_list

                )

                if not combos:

                        # If there are no other variants to combine, then simply

                        # combine all manifests.

                        combos = [()]

                # Perform the variant merge for each unique combination of

                # remaining variants.  For example, a pkgmerge of:

                #   -s arch=sparc,debug=false,flavor=32,...

                #   -s arch=sparc,debug=false,flavor=64,...

                #   -s arch=sparc,debug=true,flavor=32,...

                #   -s arch=sparc,debug=true,flavor=64,...

                #   -s arch=i386,debug=false,flavor=32,...

                #   -s arch=i386,debug=false,flavor=64,...

                #   -s arch=i386,debug=true,flavor=32,...

                #   -s arch=i386,debug=true,flavor=64,...

                #

                # ...would produce the following combinations for each variant:

                #   variant.arch

                #     debug=false, flavor=32

                #     debug=false, flavor=64

                #     debug=true, flavor=32

                #     debug=true, flavor=64

                #   variant.debug

                #     flavor=32

                #     flavor=64

                #   variant.flavor

                #

                for combo in combos:

                        # Build the list of sources, fmris, and variant values

                        # involved in this particular combo merge.

                        slist = []

                        flist = []

                        vlist = []

                        sindex = []

                        new_fmri = None

                        for j, src in enumerate(source_list):

                                if combo:

                                        # If filtering on a specific combination

                                        # then skip this source if any of the

                                        # combination parameters don't match.

                                        skip = False

                                        for vname, vval in combo:

                                                if variant_list[j].get(vname,

                                                    None) != vval:

                                                        skip = True

                                                        break