#!/usr/bin/python3 -u
# pylint: disable=E1136
"""
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.

Dependencies:
 - python3-koji, koji-utils,  krb5-workstation and python3-koji-cli-plugins
 - 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 shutil
import subprocess
import sys
import tempfile

cosa_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, f"{cosa_dir}/cosalib")
sys.path.insert(0, cosa_dir)

from cosalib.cmdlib import get_basearch, load_json, write_json

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:
    with open("/etc/redhat-release", "r") as f:
        HOST_OS = str(f.read()).strip()
except IOError:
    HOST_OS = "unknown"

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, arch specific file, 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 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, **kwargs):
        # add a dummy platform to satisfy _Build; we're not actually building a
        # new kind of image here, but _Build is geared towards that and has a
        # image_name_base() member that wants a platform string to name stuff
        self.platform = "koji"
        # Use a tempdir in builds/ because we want to use large scratch space
        self._tmpdir = tempfile.mkdtemp(prefix="koji-build", dir="builds/")
        self._state_file_tpl = self._tmpdir + "/cosa-cmd-koji-upload-{name}-{version}-{release}"
        kwargs.update({
            "require_commit": True,
            "require_cosa": True,
        })
        _Build.__init__(self, **kwargs)

    def __del__(self):
        try:
            shutil.rmtree(self._tmpdir)
        except Exception as e:
            raise Exception(
                f"failed to remove temporary directory: {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 str
        """

        rename_raw = [
            "ostree-commit"
            f'-kernel-{self.basearch}',
            f'installer-initramfs.{self.basearch}.img',
        ]

        for ending in rename_raw:
            if fname.endswith(ending):
                lname = f"{fname}.raw"
                if not os.path.exists(lname):
                    try:
                        os.symlink(fname, lname)
                    except Exception as e:
                        raise Exception("failed create symlink for {lname}", e)

                return lname, 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 = f"{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((f"using {compressor.__name__} module to mutate",
                              f"{base_name} to {new_path}"))
                    shutil.copyfileobj(infile, outfile)
                except Exception as e:
                    raise Exception(
                        f"failed to decompress file {fname} to {new_path}", e)

                finally:
                    infile.close()
                    outfile.close()

                return new_path, True

        return fname, True

    def get_rpm_list(self, host=None):
        """
        Translate commitmeta.json/HOST OS rpms into a json list
        Returns the json rpms list
        """
        components = []
        if host is None:
            rpms = self.commit["rpmostree.rpmdb.pkglist"]
        else:
            host_rpms = subprocess.check_output('rpm -qa --qf="%{NAME}:%{EPOCH}:%{RELEASE}:%{VERSION}:%{ARCH}:%{SIGMD5}:%{SIGPGP} \n"', shell=True).strip()
            rpms = (host_rpms.decode('utf-8')).split("\n")

        for rpm in rpms:
            if host is None:
                name, epoch, version, release, arch = rpm
                sigmd5, sigpgp, epoch = None, None, None
            else:
                name, epoch, release, version, arch, sigmd5, sigpgp = rpm.split(':')
            entry = {
                "type": "rpm",
                "name": name,
                "version": version,
                "release": release,
                "epoch": epoch,
                "arch": arch,
                "sigmd5": sigmd5,
                "sigpgp": sigpgp}
            components.append(entry)
        return components

    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, str)

        The return is tuple is used if the file is found and a message.
        """
        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, _, arch, _) = KOJI_CG_TYPES.get(
            check_extension, [None, None, None, None])

        if arch == "arch":
            farch = os.path.abspath(fname).split('/')[-2]
            if farch != self.basearch:
                return (False,
                        f'file arch {farch} mismatches build arch')

        if found:
            return (True, f'{check_extension} is supported by Koji')

        return (False, f'{check_extension} is not supported by Koji')

    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(f" * EXCLUDING symlink {ffile}")
                log.debug(f" *    target {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(f"Considering file file {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(f" * using {os.path.basename(lpath)} for upload name")

            # and check that a file should be uploaded
            upload_path = os.path.basename(lpath)
            supported, msg = self.supported_upload(lpath)
            if not supported:
                log.debug(f" * EXCLUDING: {msg}")
                continue

            log.debug(" * calculating checksum")
            self._found_files[lpath] = {
                "local_path": lpath,
                "upload_path": upload_path,
                "md5": md5sum_file(lpath),
                "size": os.stat(lpath).st_size
            }
            log.debug(f" * size is {self._found_files[lpath]['size']}")
            log.debug(f" * md5 is {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='[%(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(f"using kerberos auth via {keytab}")
    try:
        _ = subprocess.check_output([
            "kinit", "-f", "-t", keytab, "-k", principle])
        klist_out = subprocess.check_output(["klist", "-a"])
        log.debug(f"authenticated: \n{(klist_out.decode('utf-8'))}")
    except Exception as err:
        raise Exception("failed to auth: ", err)


class _KojiBase():
    """
    Base class for classes requiring the ability to perform operations
    in Koji.
    """

    def __init__(self, profile):
        """
        Creates a new instance of the _KojiBase.

        :param profile: Koji profile name in /etc/koji.conf.d
        :type str
        """
        self._session = None
        self._profile = profile

    @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(f"failed to authenticate to koji: {e}")

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

        self._session = session
        return session


class Search(_KojiBase):
    """
    Search for builds
    """

    def __init__(self, profile):
        """
        Creates a new instance for search.

        :param profile: Koji profile name in /etc/koji.conf.d
        :type str
        """
        super().__init__(profile)

    def get_state(self, nvr):
        """
        Return the build state.
        :param nvr: The nvr name from Brew

        For more about build state see:
          https://pagure.io/koji/blob/master/f/www/kojiweb/builds.chtml#_27
          https://pagure.io/koji/blob/master/f/tests/test_cli/test_import.py#_73
        """

        info = self.session.getBuild(nvr, strict=False)
        if (info):
            return (info['state'])

        # Don't return anything else if no build was found.
        # Since the build states are describe as numbers from 0
        # to 4, let's get an empty return
        return ""

    def check_tag(self, nvr, tag):
        """
        Check if the build contains the tag
        :param nvr: The nvr name from Brew
        :param tag: The tag to be checked
        """

        tags = self.session.listTags(build=nvr)
        for build_tag in tags:
            if tag == build_tag['name']:
                return True

        return False

    def ensure_tag(self, nvr, tag):
        """
        Ensure if the build contains the tag
        :param nvr: The nvr name from Brew
        :param tag: The tag to be checked
        """

        if not self.check_tag(nvr, tag):
            log.info('Build %s was not tagged. Adding tag: %s' % (nvr, tag))
            task_id = self.session.tagBuild(tag, nvr)
            task_result = klib.watch_tasks(self.session, [task_id], quiet=True, poll_interval=15)
            if task_result != 0:
                raise Exception('failed to tag builds')
            log.info('Tag %s successfully added' % (tag))
        else:
            log.info('Tag %s already exists' % (tag))


class Reserve(_KojiBase):
    """
    Reserves a place in Koji for later archival.
    """

    def __init__(self, profile):
        """
        Creates a new instance for reservation.

        :param profile: Koji profile name in /etc/koji.conf.d
        :type str
        """
        super().__init__(profile)

    def reserve_id(self, build, create_file=None):
        """
        Reserves a koji id for use during upload. This must be called
        after build a meta.json is available for review.

        See: https://docs.pagure.org/koji/content_generators/#api
        """
        log.info("Reserving a unique koji id")

        # The koji/brew NVR is constructed like so:
        # Name = "rhcos-$arch", like `rhcos-x86_64`
        # Version = Everything before `-` in RHCOS version
        # Release = Everything after `-` in RHCOS version
        #
        # Example: RHCOS Build ID: 414.92.202307170903-0 for x86_64
        #   Name = rhcos-x86_64
        #   Version = 414.92.202307170903
        #   Release = 0
        #   NVR = rhcos-x86_64-414.92.202307170903-0
        version, release = build.build_id.split('-')
        data = {
            "name": f"{build.build_name}-{build.basearch}",
            "release": release,
            "version": version,
            "cg": "coreos-assembler",
        }

        log.debug(f"reserve payload {data}")
        koji_reservation = self.session.CGInitBuild(data['cg'], data)
        log.info(f"Koji reservation id created: {koji_reservation['build_id']}")

        if create_file is not None:
            state_file = build._state_file_tpl.format(
                name=data['name'],
                version=data['version'],
                release=data['release'],
                build_id=koji_reservation['build_id'],
                token=koji_reservation['token'])
            write_json(state_file, (data, koji_reservation))
            log.info(f"reserve data written to {state_file}")
            return state_file

        build.meta['koji'] = {
            'build_id': koji_reservation['build_id'],
            'token': koji_reservation['token'],
            'release': release
        }
        build.meta_write()


class Upload(_KojiBase):
    """ 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.
        """
        super().__init__(profile)
        self._build = in_build
        self._build_id = None
        self._manifest = None
        self._owner = owner.split('@')[0]
        self._token = None
        self._remote_directory = None
        self._session = None
        self._tag = tag
        self._image_files = None
        self._reserve_id_file = None
        self._retry_attempts = 2
        self._uploaded = False
        self._s3 = None
        self._s3_bucket = None
        self._s3_key = None
        self._s3_url = None

        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

    def get_file_meta(self, 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 = obj.get("upload_path").split('.')[-2]
            ext = f"{sub_ext}.{ext}"

        log.debug(f"File {(obj.get('upload_path'))} is type '{ext}': {obj}")
        (description, ext, arch, etype) = KOJI_CG_TYPES.get(
            ext, [None, None, None, None])

        if arch != "noarch":
            arch = self.build.basearch

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

        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 = self.get_file_meta(value)
            if file_output is not None:
                if "commitmeta.json" in value['upload_path']:
                    file_output["components"] = self.build.get_rpm_list()
                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

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

        """
        Koji has a couple of checks to ensure the reservation data (build_Id, release, name
        and version) that is passed in the build import matchs with the build data.
        For our case, the release data is done using time (“%H%M%S”), it always to be a mismatch
        with the release time created for the build.
        Let's just replace the build_id and the release here, and let the further checks for Koji.
        """
        if self._reserve_id_file is not None:
            if os.path.isfile(self._reserve_id_file):
                log.debug(f"loading reservation data from state file {self._reserve_id_file}")
                data = load_json(self._reserve_id_file)
                self._build_id = data[1]['build_id']
                self._release = data[0]['release']
                self._token = data[1]["token"]

                self.build.meta['koji'] = {
                    'build_id': self._build_id
                }
                self.build.meta_write()
        else:
            try:
                if self.build.meta['koji']['build_id'] is not None:
                    log.debug(f"loading reservation data from {self.build.meta['koji']}")
                    self._build_id = self.build.meta['koji']['build_id']
                    self._release = self.build.meta['koji']['release']
                    self._token = self.build.meta['koji']['token']
            except:
                pass

        if self._s3 is not None:
            self.build.meta['s3'] = {
                'bucket': self._s3_bucket,
                'key': self._s3_key,
                'public-url': self._s3_url
            }
            self.build.meta_write()

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

        log.debug(f"Preparing manifest for {(len(self.image_files))} files")
        # The koji/brew NVR is constructed like so:
        # Name = "rhcos-$arch", like `rhcos-x86_64`
        # Version = Everything before `-` in RHCOS version
        # Release = Everything after `-` in RHCOS version
        #
        # Example: RHCOS Build ID: 414.92.202307170903-0 for x86_64
        #   Name = rhcos-x86_64
        #   Version = 414.92.202307170903
        #   Release = 0
        #   NVR = rhcos-x86_64-414.92.202307170903-0
        version, release = self.build.build_id.split('-')
        self._manifest = {
            "metadata_version": 0,
            "build": {
                "build_id": self._build_id,
                "end_time": stamp,
                "extra": {
                    "typeinfo": {
                        "image": {
                            "arch": self.build.basearch,
                        }
                    }
                },
                "name": f"{self.build.build_name}-{self.build.basearch}",
                "release": release,
                "owner": self._owner,
                "source": source['origin'],
                "start_time": stamp,
                "version": version
            },
            "buildroots": [{
                "id": 1,
                "host": {
                    "os": HOST_OS,
                    "arch": self.build.basearch
                },
                "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.basearch,
                    "name": "coreos-assembler"
                },
                "components": self.build.get_rpm_list('host'),
                "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

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

    def tag_build(self, nvr, tag):
        """ Tag build in Koji """
        log.info('tagging %s into %s' % (nvr, tag))
        task_id = self.session.tagBuild(tag, nvr)
        task_result = klib.watch_tasks(self.session, [task_id], poll_interval=15)
        if task_result != 0:
            raise Exception('failed to tag builds')

    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 = klib.unique_path(f"{self.build.build_id}-cosa")
        log.debug(f"uploading files to {serverdir}")
        meta_path = ""
        meta_remote_path = ""
        callback = klib._progress_callback

        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(f"Uploading {local_path} to {serverdir}/{remote_path}")
            if ("meta.json" in local_path) and ("commitmeta.json" not in local_path):
                # Keep location for meta.json
                meta_path = local_path
                meta_remote_path = remote_path
            self.session.uploadWrapper(local_path, serverdir, name=remote_path,
                                       callback=callback)
            if callback:
                pass

        self._uploaded = True
        self._remote_directory = serverdir
        cginfo = ""

        # If the final Import fails retry it
        for attempt in range(1, self._retry_attempts + 1):
            self.manifest
            try:
                log.info(f"Finishing Koji Upload - Try {attempt}/{self._retry_attempts}")
                cginfo = self.session.CGImport(self.manifest, serverdir, token=self._token)
                log.info(json.dumps(cginfo, sort_keys=True, indent=3))
                # Koji has some flakes here, confirm the nvr to make sure the upload worked
                if cginfo['nvr']:
                    log.info(f"Recorded build {cginfo['nvr']}")
                    break
            except Exception as e:
                if attempt < self._retry_attempts:
                    if (self._token):
                        log.info(f"Koji upload failed with: {e}. Canceling koji reservation \
                        {self._build_id} and retrying")
                        self.session.CGRefundBuild("coreos-assembler", self._build_id, self._token, state=4)
                        log.info(f"Koji id reservation canceled: {self._build_id}")
                        Reserve(self._profile).reserve_id(self.build)
                        # Force manifest and checksum/size update for the new meta generated
                        self._manifest = None
                        self._image_files = None
                        self.build._found_files[meta_path]['md5'] = md5sum_file(meta_path)
                        self.build._found_files[meta_path]['size'] = os.stat(meta_path).st_size
                        # Update the meta.json with the new id/token into serverdir
                        self.session.uploadWrapper(meta_path, serverdir, name=meta_remote_path,
                                                   callback=callback)
                        if callback:
                            pass
                else:
                    raise Exception("Koji upload failed with:", e)

        if self.session.hasPerm('tag'):
            self.tag_build(cginfo['nvr'], self._tag)
        else:
            log.info("Principle does not have tagging permissions.  Not tagging build")

        return cginfo


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

    parser = argparse.ArgumentParser()
    parent_parser = argparse.ArgumentParser(
        prog="CoreOS Assembler Koji Uploader",
        description='Archive build artifacts, logs, and metadata in Koji.',
        usage="""
Reserve an id and 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.

Examples:
    $ cmd-koji-upload reserve-id \
        --keytab keytab \
        --profile koji \
        --name rhcos \

    $ cmd-koji-upload upload \
        --reserve-id \
        --keytab keytab \
        --owner me@FEDORA.COM \
        --profile koji
    $ cmd-koji-upload search \
        --nvr nvr \
        --keytab keytab \
        --owner me@FEDORA.COM \
        --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"""
    )

    parent_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.
    parent_parser.add_argument("--build", default="latest",
                        help="Override build id, defaults to latest")
    parent_parser.add_argument("--buildroot", default="builds",
                        help="Build diretory")
    parent_parser.add_argument("--dump", default=False, action='store_true',
                        help="Dump the manifest and exit")
    parent_parser.add_argument("--no-upload", default=False, action='store_true',
                        help="Do not upload, just parse the build")
    parent_parser.add_argument("--arch", default=get_basearch(),
                        help="Set the build architecture")

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

    sub_commands = parser.add_subparsers(
        dest="_command",
        help="Other commands")

    upload_cmd = sub_commands.add_parser(
        "upload", help="Uploads to koji", parents=[parent_parser], add_help=False)

    search_cmd = sub_commands.add_parser(
        "search", help="Search for a build", parents=[parent_parser], add_help=False)

    ensure_cmd = sub_commands.add_parser(
        "ensure-tag", help="Ensure the build tag is correct", parents=[parent_parser], add_help=False)
    sub_commands.add_parser(
        "reserve-id", help="Reserves a koji id", parents=[parent_parser], add_help=False)

    upload_cmd.add_argument(
        '--tag', required=True,
        default=os.environ.get("KOJI_TAG", None),
        help='tag this build, eg. awesome-candidate')

    upload_cmd.add_argument(
        '--retry-attempts', required=False,
        help='If the upload fails retry it n times')

    upload_cmd.add_argument(
        '--reserve-id', required=False, action='store_true',
        help='Creates a reservation ID in the upload process')

    upload_cmd.add_argument(
        '--reserve-id-state-file', required=False,
        help='Uses the path for a reservation file previous created')

    upload_cmd.add_argument(
        '--s3-bucket', required=False,
        help='Store bucket information in meta.json')

    upload_cmd.add_argument(
        '--s3-key', required=False,
        help='Store key information in meta.json')

    upload_cmd.add_argument(
        '--s3-url', required=False,
        help='Store url information in meta.json')

    search_cmd.add_argument(
        '--nvr',  required=True,
        help='NVR to look for')

    ensure_cmd.add_argument(
        '--nvr',  required=True,
        help='NVR to look for')

    ensure_cmd.add_argument(
        '--tag',  required=True,
        help='Ensure the tag if the build does not have it')

    args, extra_args = parser.parse_known_args()
    set_logger(args.log_level)

    if args._command is None:
        sys.argv.insert(3, 'upload')
        log.warning(
            f'Calling {sys.argv[0]} a subcommand is deprecated. Please update'
            ' your command call with a subcommand. Defaulting to "upload"'
            ' for subcommandless run.')

    args = parser.parse_args(namespace=args)

    build = Build(buildroot=args.buildroot, build=args.build, arch=args.arch)

    if args.auth:
        kinit(args.keytab, args.owner)
    if args._command == 'search':
        print(Search(args.profile).get_state(args.nvr))
    if args._command == 'ensure-tag':
        Search(args.profile).ensure_tag(args.nvr, args.tag)
    if args._command == 'upload':

        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)
        if args.retry_attempts:
            upload._retry_attempts = int(args.retry_attempts)
        if args.reserve_id:
            Reserve(args.profile).reserve_id(build)
        if args.reserve_id_state_file:
            upload._reserve_id_file = args.reserve_id_state_file
        if args.s3_bucket or args.s3_key or args.s3_url:
            upload._s3 = True
            upload._s3_bucket = args.s3_bucket
            upload._s3_key = args.s3_key
            upload._s3_url = args.s3_url

        build.build_artifacts()
        upload.upload()
    elif args._command == 'reserve-id':
        return Reserve(args.profile).reserve_id(build, "create_file")


if __name__ == '__main__':
    cli()
