cmd-koji-upload

#!/usr/bin/python -u
# pylint: disable=C0103

"""
cmd-koji-upload performs the required steps to make COSA a Koji Content
Generator. When running this in an automated fashion, you will need a Kerberos
keytab file for a service account.

Python2: due to an issue with Kerberos Auth issue which returns errors like:
[auth_gssapi:error] NO AUTH DATA Client did not send any authentication headers
[auth_gssapi:error] GSS ERROR gss_localname() failed:...
We are unable to use Python3. After inquiring with the Koji team, we were
advised to use Python2. This code should port rather cleanly to Python3 once
the Kerberos issue is cleared.

Dependencies:
 - python-koji, krb5-workstation and python-krbV
 - Kerberos credentials against Koji _in_ keytab format
 - content generator permissions for your user or service account

See cli() for usage information.
"""
import argparse
import datetime
import gzip
import hashlib
import json
import koji
import koji_cli.lib as klib
import logging as log
import os.path
import platform
import shutil
import subprocess
import sys
import tempfile

cosa_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ("%s/cosalib" % cosa_dir))
sys.path.insert(0, cosa_dir)

try:
    from cosalib.build import _Build
except ImportError:
    # We're in the container and can't sense the cosa path
    sys.path.insert(0, '/usr/lib/coreos-assembler/cosalib')
    from cosalib.build import _Build

try:
    # Available in Koji v1.17, https://pagure.io/koji/issue/975
    from koji_cli.lib import unique_path
    from koji_cli.lib import progress_callback as cb
except ImportError:
    from koji_cli.lib import _unique_path as unique_path
    from koji_cli.lib import _progress_callback as cb

try:
    with open("/etc/redhat-release", "r") as f:
        HOST_OS = str(f.read()).strip()
except IOError:
    HOST_OS = "unknown"

# ARCH is the current machine architecture
ARCH = platform.machine()

COSA_INPATH = "/cosa"

# Content generator types as defined in:
# https://pagure.io/koji/blob/master/f/docs/schema.sql
#   DO NOT add arbitrary types or extensions here. If the archive type
#   is unknown to Koji, then the upload will fail.
KOJI_CG_TYPES = {
    # key: (description, extension, architecture, type)
    "iso": ("CD/DVD Image", "iso", ARCH, "image"),
    "json": ("JSON data", "json", "noarch", "source"),
    "log": ("log file", "log", "noarch", "log"),
    "ova": ("Open Virtualization Archive", "ova", ARCH, "image"),
    "qcow": ("QCOW image", "qcow", ARCH, "image"),
    "qcow2": ("QCOW2 image", "qcow2", ARCH, "image"),
    "qcow2.gz": ("Compressed QCOW2", "qcow2.gz", ARCH, "image"),
    "qcow2.xz": ("Compressed QCOW2", "qcow2.gz", ARCH, "image"),
    "raw": ("Raw disk image", "raw", ARCH, "image"),
    "raw.gz": ("Compressed Raw", "raw.gz", ARCH, "image"),
    "raw.xz": ("Compressed Raw", "raw.xz", ARCH, "image"),
    "tar": ("Tar file", "tar", "noarch", "source"),
    "tar.gz": ("Tar file", "tar.gz", "noarch", "source"),
    "tar.xz": ("Tar file", "tar.xz", "noarch", "source"),
    "vdi": ("VirtualBox Virtual Disk Image", "vdi", ARCH, "image"),
    "vhd": ("Hyper-V image", "vhd", ARCH, "image"),
    "vhdx": ("Hyper-V Virtual Hard Disk v2 image", "vhdx", ARCH, "image"),
    "vmdk": ("vSphere image", "vmdk", ARCH, "image"),
    "yaml": ("YAML data", "yaml", "noarch", "source")
}

# These artifacts need to be renamed on upload. Koji file extension matching
# against content types.
RENAME_RAW = ['initramfs.img', '-kernel', "ostree-commit"]

# These are compressed extensions that are used to determine if name managling
# might be needed.
COMPRESSION_EXT = ["gz", "xz"]


def md5sum_file(path):
    """
    Calculates the md5 sum from a path.
    Py3 TODO: use md5sum_file in cmdlib.py, when porting to Python3.

    :param path: file name to checksum
    :returns str

    Returns the hexdigest of the file.
    """
    h = hashlib.md5()
    with open(path, 'rb', buffering=0) as data:
        for b in iter(lambda: data.read(128 * 1024), b''):
            h.update(b)
    return h.hexdigest()


class Build(_Build):
    """
    Koji implementation of Build.
    """
    def __init__(self, *args, **kwargs):
        self._tmpdir = tempfile.mkdtemp(prefix="koji-build")
        _Build.__init__(self, *args, **kwargs)

    def __del__(self):
        try:
            shutil.rmtree(self._tmpdir)
        except Exception as e:
            raise Exception("failed to remove temporary directory: %s",
                            self._tmpdir, e)

    def rename_mutator(self, fname):
        """
        If a file needs to be renamed because of a Koji rule, rename
        the file to a `.raw`. For some types, its technically incorrect.

        :param fname: file name to check if it needs to be renamed
        :type str
        :return srt
        """
        for ending in RENAME_RAW:
            if fname.endswith(ending):
                return "%s.raw", True
        return fname, None

    def decompress_mutator(self, fname):
        """
        Calculate the mutated name and return a file object suitable
        for reading the file. If the file is supported as is by Koji,
        the original file is returned.

        :param fname: name of the local file
        :type str
        :return str
        :return file
        """
        for x in COMPRESSION_EXT:
            if not fname.endswith(x):
                continue

            base_name = os.path.basename(fname)
            base = os.path.splitext(base_name)[0]
            ftype = (os.path.splitext(base)[1]).replace(".", "")
            ctype = "%s.%s" % (ftype, x)
            if ctype not in KOJI_CG_TYPES and ftype in KOJI_CG_TYPES:
                if x == "gz":
                    compressor = gzip
                elif x == "xz":
                    raise Exception("not supported yet")

                new_path = os.path.join(self._tmpdir, base)
                try:
                    infile = compressor.open(fname)
                    outfile = open(new_path, 'wb+')
                    log.info("using %s module to mutate %s to %s",
                             compressor.__name__, base_name, new_path)
                    shutil.copyfileobj(infile, outfile)
                except Exception as e:
                    raise Exception("failed to decompress file %s to %s: %s",
                                    fname, new_path, e)
                finally:
                    infile.close()
                    outfile.close()

                return new_path, True

        return fname, True

    def mutate_for_koji(self, fname):
        """
        Koji is _so_ pendantic about the naming of files and their extensions,
        such that "vhd.gz" is not allowed, but "vhd" is. In the event that a
        file needs to be mutated, this function will do that.

        :param fname: name of the file to mutate
        :type str

        :Returns: location of fname or the name of the mutated file
        """
        (new_name, _) = self.decompress_mutator(fname)
        (new_name, _) = self.rename_mutator(new_name)
        if fname != new_name:
            return new_name
        return fname

    def supported_upload(self, fname):
        """
        Helper to return if a file should be uploaded
        :param fname: name of file to check against Koji table for uploading.
        :type str
        :returns bool

        Returns true if the file is known to Koji.
        """
        base = os.path.basename(fname)
        check_extension = base.split(".")[-1]
        for extension in COMPRESSION_EXT:
            if fname.endswith(extension):
                check_extension = base.split(".")[-2]

        found = KOJI_CG_TYPES.get(check_extension)
        if found:
            return True
        return False

    def _build_artifacts(self, *args, **kwargs):
        """
        Implements the building of artifacts. Walk the build root and
        prepare a list of files in it. While commitmeta.json provides
        these artifacts, the information about the individual files is
        not compatible with the Koji upload (e.g. SHA, type has to be
        translated, strict filename conventions)

        :param args: All non-keyword arguments
        :type args: list
        :param kwargs: All keyword arguments
        :type kwargs: dict
        """
        # locate all the build artifacts in the build directory.
        files = []
        for ffile in os.listdir(self.build_dir):
            files.append(os.path.join(self.build_dir, ffile))
            if os.path.islink(ffile):
                log.debug(" * EXCLUDING symlink '%s'", ffile)
                log.debug(" *    target '%s'", os.path.realpath(ffile))
                continue

        # add the coreos assembler files
        for ffile in os.listdir(COSA_INPATH):
            files.append(os.path.join(COSA_INPATH, ffile))

        # process the files that were found
        for ffile in files:
            lpath = os.path.abspath(ffile)
            log.debug("Considering file file '%s'", lpath)

            # if the file is mutated (renamed, uncompressed, etc)
            # we want to use that file name instead
            mutated_path = self.mutate_for_koji(lpath)
            if mutated_path != lpath:
                lpath = mutated_path
            log.debug(" * using %s for upload name", lpath)

            # and check that a file should be uploaded
            upload_path = os.path.basename(lpath)
            if not self.supported_upload(lpath):
                log.debug(" * EXCLUDING file '%s'", ffile)
                log.debug("   File type is not supported by Koji")
                continue

            # os.path.getsize uses 1kB instead of 1KB. So we use stat instead.
            fsize = subprocess.check_output(["stat", "--format", '%s', lpath])
            log.debug(" * calculating checksum")
            self._found_files[lpath] = {
                "local_path": lpath,
                "upload_path": upload_path,
                "md5": md5sum_file(lpath),
                "size": int(fsize)
            }
            log.debug(" * size is %s", self._found_files[lpath]["size"])
            log.debug(" * md5 is %s", self._found_files[lpath]["md5"])


def set_logger(level):
    """
    Set the log level
    :param level: set the log level
    :type str
    """
    sl = log.DEBUG
    if level == "error":
        sl = log.ERROR
    elif level == "warn":
        sl = log.WARN
    elif level == "info":
        sl = log.INFO

    log.basicConfig(format='[%(asctime)s  %(levelname)s]: %(message)s',
                    level=sl)


def kinit(keytab, principle):
    """
    Execute the Kerberos Auth using the provided keytab:
    :param keytab: file name of the keytab used for Kerberos Authentication
    :type str
    :param principle: name of the keytab's principle
    :type str
    """
    if keytab is None:
        raise Exception("keytab file is not defined")

    log.info("using kerberos auth via %s", keytab)
    try:
        _ = subprocess.check_output([
            "kinit", "-f", "-t", keytab, "-k", principle])
        klist_out = subprocess.check_output(["klist", "-a"])
        log.debug("authenticated: \n%s", klist_out.decode("utf-8"))
    except Exception as err:
        raise Exception("failed to auth: ", err)


class Upload():
    """ Upload generates the manifest for a build and uploads to a
        Koji Server. Upload treats each instance as a separate build; multiple
        innovations of the Upload should be separate instances.
    """

    def __init__(self, in_build, owner, tag, profile, login_now=False):
        """ Initialize an instance of the Upload object.
        :param in_build: build object to process for uploading
        :type Build
        :param owner: name of the upload owner
        :type str
        :param tag: name of the tag to apply to the upload
        :type str
        :param profile: Koji profile name in /etc/koji.conf.d
        :type str
        :param login_now: authenticate via kerberos immediately
        :type bool
        :raises Exception

        During the init of the Upload object, we check to make sure that an
        upload is likely to be succeed. We want to fail early in-case there is
        missing required information.
        """
        self._build = in_build
        self._manifest = None
        self._owner = owner.split('@')[0]
        self._profile = profile
        self._remote_directory = None
        self._session = None
        self._tag = tag
        self._image_files = None
        self._uploaded = False

        if self._tag is None:
            raise Exception("build tag must be set")
        if self._owner is None:
            raise Exception("owner must be set")
        if self.build is None:
            raise Exception("Build object must be provided")
        if self._profile is None:
            raise Exception("profile must not be None")

        if login_now:
            try:
                _ = self.session
            except Exception as err:
                raise Exception("failed to login into Koji instance", err)
            self.verify_tag(self._tag)

    @property
    def build(self):
        """ get the Build that Upload will act on """
        return self._build

    @staticmethod
    def get_file_meta(obj):
        """
        Generate the required dict for specific types of files.
        @param dict: meta-data for parsing

        HERE BE DRAGONS: YOU MUST TEST ANY CHANGES AGAINST KOJI. CHANGING
        THIS STRUCTURE CAN AND WILL BREAK UPLOADS. THE STRUCTURE FOR A
        CONTENT GENERATOR IS NOT DOCUMENTED.
        """
        if not isinstance(obj, dict):
            raise Exception("cannot parse file meta-data, invalid type")

        ext = os.path.splitext(obj.get("upload_path"))[-1]
        ext = ext.lstrip('.')
        if ext in COMPRESSION_EXT:
            # find sub extension, e.g. "tar" in "tar.gz"
            sub_ext = os.path.splitext(obj.get("upload_path"))[0].lstrip('.')
            ext = "%s.%s" % (sub_ext, ext)

        log.debug("File %s should be of type %s: %s ", obj.get("upload_path"),
                  ext, obj)
        (description, ext, arch, etype) = KOJI_CG_TYPES.get(
            ext, [None, None, None, None])

        file_meta = {
            "arch": arch,
            "buildroot_id": 1,
            "checksum": obj["md5"],
            "checksum_type": "md5",
            "filename": obj['upload_path'],
            "filesize": obj["size"],
            "type": ext,
            "extra": {"image": {"arch": arch}}
        }

        if etype not in ["image", "source"]:
            del file_meta['extra']
            file_meta['type'] = etype

        if description is None:
            return None

        return file_meta

    @property
    def image_files(self):
        """ Generate outputs prepares the output listing. """
        if self._image_files is not None:
            return self._image_files

        outputs = []
        for _, value in (self.build).get_artifacts():
            file_output = Upload.get_file_meta(value)
            if file_output is not None:
                outputs.append(file_output)
        self._image_files = outputs
        return self._image_files

    @property
    def manifest(self):
        """
        Generate the core json used to tell Koji what the build looks like
        """
        if self._manifest is not None:
            return self._manifest

        source = self.build.get_meta_key(
            "meta", self.build.ckey("container-config-git"))

        now = datetime.datetime.utcnow()
        stamp = now.strftime("%s")

        log.debug("Preparing manifest for %s files", len(self.image_files))
        self._manifest = {
            "metadata_version": 0,
            "build": {
                "end_time": stamp,
                "extra": {
                    "typeinfo": {
                        "image": {
                            "arch": ARCH
                        }
                    }
                },
                "name": self.build.build_name,
                "release": now.strftime("%H%M%S"),
                "owner": self._owner,
                "source": source['origin'],
                "start_time": stamp,
                "version": "%s" % self.build.build_id
            },
            "buildroots": [{
                "id": 1,
                "host": {
                    "os": HOST_OS,
                    "arch": self.build.arch
                },
                "content_generator": {
                    "name": "coreos-assembler",
                    "version": self.build.get_sub_obj(
                        "meta",
                        self.build.ckey("container-config-git"), "commit")
                },
                "container": {
                    "type": "docker",
                    "arch": self.build.arch,
                    "name": "coreos-assembler"
                },
                "components": "",
                "extra": {
                    "coreos-assembler": {
                        "build_id": 1,
                        "builder_image_id": 1
                    }
                },
                "type": "image",
                "tools": [
                    {
                        "name": "coreos-assembler",
                        "version": self.build.get_sub_obj(
                            "meta",
                            self.build.ckey("container-config-git"), "commit")
                    }
                ]
            }],
            "output": self.image_files
        }
        return self._manifest

    @property
    def session(self):
        """
        Return an authenticated Koji session
        """

        if self._session is not None:
            return self._session

        mykoji = koji.get_profile_module(self._profile)
        opts = mykoji.grab_session_options(mykoji.config)
        session = mykoji.ClientSession(mykoji.config.server, opts)

        try:
            klib.activate_session(session, mykoji.config)
            assert session.logged_in
            log.info("logged into koji server")
        except Exception as e:
            raise Exception("failed to authenticate to koji: %s" % e)

        if session is None:
            raise Exception("failed to get session from koji server")

        self._session = session
        return session

    def verify_tag(self, tag):
        """ Verify that a tag exists in this Koji instance """
        taginfo = self._session.getTag(tag)
        if not taginfo:
            raise RuntimeError('tag %s is not present in Koji' % tag)

    def upload(self):
        """
        Upload all files to a remote directory in Koji. Content generators
        upload all the artifacts first and then imports them based on the
        manifest output.

        The Koji Python lib maintains the filename in the upload; there is no
        option to upload with a different name. If the name needs to be changed
        a symlink is created in temporary directory and is used as the upload
        file.
        """
        serverdir = unique_path("%s-cosa" % self.build.build_id)
        callback = cb

        log.debug('uploading files to %s', serverdir)
        for _, meta in (self.build).get_artifacts():
            local_path = meta['local_path']    # the local file to upload
            remote_path = meta['upload_path']  # the name of the file to upload
            log.info("Uploading %s to %s/%s", local_path, serverdir,
                     remote_path)
            self.session.uploadWrapper(local_path, serverdir, name=remote_path,
                                       callback=callback)
            if callback:
                print('')

        self._uploaded = True
        self._remote_directory = serverdir
        cginfo = self.session.CGImport(self.manifest, serverdir)
        log.info(json.dumps(cginfo, sort_keys=True, indent=3))
        log.info("recorded build %s", cginfo['nvr'])
        return cginfo


def cli():
    """ cli implements command-line innovation """

    parser = argparse.ArgumentParser(
        prog="CoreOS Assembler Koji Uploader",
        description='Archive build artifacts, logs, and metadata in Koji.',
        usage="""
Upload a CoreOS Assembler (COSA) created Build to a Koji Server.

Note: the typical use case for this program is in an automated fashion, and
running from within the COSA container yourself. Unless you are debugging COSA,
this is likely not what are looking for.

To use this program, you will need:
    1) An account on the Koji Server
    2) A Keytab file with your Kerberos Credentials
    3) A completed build.

Example:
    $ cmd-koji-upload \
        --build_root=/src/build \
        --keytab keytab \
        --owner me@FEDORA.COM \
        --tag rhaos-4.1-rhel-8-build \
        --profile koji

Environment variables are supported:
    - KOJI_USERNAME will set the owner
    - KOJI_TAG will set the tag for the build
    - KOJI_PROFILE is the configuration in /etc/koji.confi.d that will be used
    - KEYTAB will set the location for the keytab file"""
    )

    parser.add_argument("--log-level",
                        default=os.environ.get("COSA_LOG_LEVEL", "info"),
                        choices=["warn", "error", "debug", "info"],
                        help="Set the log level")

    # Options for finding the build.
    parser.add_argument("--build", default="latest",
                        help="Override build id, defaults to latest")
    parser.add_argument("--buildroot", default="builds",
                        help="Build diretory")
    parser.add_argument("--dump", default=False, action='store_true',
                        help="Dump the manfiest and exit")
    parser.add_argument("--no-upload", default=False, action='store_true',
                        help="Do not upload, just parse the build")

    # Koji specific options
    parser.add_argument("--no-auth", action='store_false', dest="auth",
                        help="Skip Kerberos auth, use if already auth'd")
    parser.add_argument("--keytab",
                        default=os.environ.get("KOJI_KEYTAB", None),
                        help="location of the keytab file to use for auth")
    parser.add_argument('--owner', required=True,
                        default=os.environ.get("KOJI_USERNAME", None),
                        help='koji user name that owns this build')
    parser.add_argument('--tag', required=True,
                        default=os.environ.get("KOJI_TAG", None),
                        help='tag this build, eg. awesome-candidate')
    parser.add_argument('--profile', required=True,
                        default=os.environ.get("KOJI_PROFILE", None),
                        help='profile to use, e.g. prod, stage, test')
    args = parser.parse_args()

    set_logger(args.log_level)

    build = Build(args.buildroot, args.build)
    if args.auth:
        kinit(args.keytab, args.owner)

    build.build_artifacts()
    upload = Upload(build, args.owner, args.tag, args.profile)
    if args.dump:
        print(json.dumps(upload.manifest, sort_keys=True, indent=3))
        return

    if args.no_upload is False:
        upload = Upload(build, args.owner, args.tag, args.profile)

    upload.upload()


if __name__ == '__main__':
    cli()