lorax

#!/usr/bin/python2
#
# lorax
#
# Copyright (C) 2009-2014 Red Hat, Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Red Hat Author(s):  Martin Gracik <mgracik@redhat.com>
#
from __future__ import print_function

import logging
log = logging.getLogger("lorax")
program_log = logging.getLogger("program")
pylorax_log = logging.getLogger("pylorax")
dnf_log = logging.getLogger("dnf")


import sys
import os
import tempfile
from optparse import OptionParser, OptionGroup
import shutil

import dnf
import pylorax

def setup_logging(opts):
    # Setup logging to console and to logfile
    log.setLevel(logging.DEBUG)
    pylorax_log.setLevel(logging.DEBUG)

    sh = logging.StreamHandler()
    sh.setLevel(logging.INFO)
    fmt = logging.Formatter("%(asctime)s: %(message)s")
    sh.setFormatter(fmt)
    log.addHandler(sh)
    pylorax_log.addHandler(sh)

    fh = logging.FileHandler(filename=opts.logfile, mode="w")
    fh.setLevel(logging.DEBUG)
    fmt = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
    fh.setFormatter(fmt)
    log.addHandler(fh)
    pylorax_log.addHandler(fh)

    # External program output log
    program_log.setLevel(logging.DEBUG)
    logfile = os.path.abspath(os.path.dirname(opts.logfile))+"/program.log"
    fh = logging.FileHandler(filename=logfile, mode="w")
    fh.setLevel(logging.DEBUG)
    program_log.addHandler(fh)

    # dnf logging
    dnf_log.setLevel(logging.DEBUG)
    logfile = os.path.abspath(os.path.dirname(opts.logfile))+"/dnf.log"
    fh = logging.FileHandler(filename=logfile, mode="w")
    fh.setLevel(logging.DEBUG)
    dnf_log.addHandler(fh)


def main(args):
    # get lorax version
    try:
        from pylorax import version
        vernum = version.num
    except ImportError:
        vernum = "devel"

    version = "{0}-{1}".format(os.path.basename(args[0]), vernum)
    usage = "%prog -p PRODUCT -v VERSION -r RELEASE -s REPOSITORY OUTPUTDIR"

    parser = OptionParser(usage=usage)

    # required arguments for image creation
    required = OptionGroup(parser, "required arguments")
    required.add_option("-p", "--product", help="product name",
                        metavar="STRING")
    required.add_option("-v", "--version", help="version identifier",
                        metavar="STRING")
    required.add_option("-r", "--release", help="release information",
                        metavar="STRING")
    required.add_option("-s", "--source",
                        help="source repository (may be listed multiple times)",
                        metavar="REPOSITORY", action="append", default=[])

    # optional arguments
    optional = OptionGroup(parser, "optional arguments")
    optional.add_option("-m", "--mirrorlist",
                        help="mirrorlist repository (may be listed multiple times)",
                        metavar="REPOSITORY", action="append", default=[])
    optional.add_option("-t", "--variant",
                        help="variant name", metavar="STRING")
    optional.add_option("-b", "--bugurl",
                        help="bug reporting URL for the product", metavar="URL",
                        default="your distribution provided bug reporting tool")
    optional.add_option("--isfinal", help="",
                        action="store_true", default=False, dest="isfinal")
    optional.add_option("-c", "--config", default="/etc/lorax/lorax.conf",
                        help="config file", metavar="STRING")
    optional.add_option("--proxy", default=None,
                        help="repo proxy url:port", metavar="STRING")
    optional.add_option("-i", "--installpkgs", default=[],
                        action="append", metavar="STRING",
                        help="package glob to install before runtime-install.tmpl runs. (may be listed multiple times)")
    optional.add_option("--buildarch", default=None,
                        help="build architecture", metavar="STRING")
    optional.add_option("--volid", default=None,
                        help="volume id", metavar="STRING")
    optional.add_option("--macboot", help="",
                        action="store_true", default=True, dest="domacboot")
    optional.add_option("--nomacboot", help="",
                        action="store_false", dest="domacboot")
    optional.add_option("--noupgrade", help="",
                        action="store_false", default=True, dest="doupgrade")
    optional.add_option("--logfile", default="./lorax.log",
                        help="Path to logfile")
    optional.add_option("--tmp", default="/var/tmp",
                        help="Top level temporary directory" )
    optional.add_option("--cachedir", default=None,
                        help="DNF cache directory. Default is a temporary dir.")
    optional.add_option("--workdir", default=None,
                        help="Work directory, overrides --tmp. Default is a temporary dir under /var/tmp")
    optional.add_option("--force", default=False, action="store_true",
                        help="Run even when the destination directory exists")
    optional.add_option("--add-template", dest="add_templates",
                        action="append", help="Additional template for runtime image",
                        default=[])
    optional.add_option("--add-template-var", dest="add_template_vars",
                        action="append", help="Set variable for runtime image template",
                        default=[])
    optional.add_option("--add-arch-template", dest="add_arch_templates",
                        action="append", help="Additional template for architecture-specific image",
                        default=[])
    optional.add_option("--add-arch-template-var", dest="add_arch_template_vars",
                        action="append", help="Set variable for architecture-specific image",
                        default=[])

    # add the option groups to the parser
    parser.add_option_group(required)
    parser.add_option_group(optional)

    # add the show version option
    parser.add_option("-V", help="show program's version number and exit",
                      action="store_true", default=False, dest="showver")

    # parse the arguments
    opts, args = parser.parse_args()

    if opts.showver:
        print(version)
        sys.exit(0)

    try:
        outputdir = os.path.abspath(args[0])
    except IndexError:
        parser.error("missing one or more required arguments")

    # check for the required arguments
    if not opts.product or not opts.version or not opts.release \
            or not opts.source or not outputdir:
        parser.error("missing one or more required arguments")

    if not opts.force and os.path.exists(outputdir):
        parser.error("output directory %s should not exist." % outputdir)

    opts.logfile = os.path.abspath(opts.logfile)
    if not os.path.exists(os.path.dirname(opts.logfile)):
        os.makedirs(os.path.dirname(opts.logfile))
    if opts.cachedir:
        opts.cachedir = os.path.abspath(opts.cachedir)
    if opts.workdir:
        opts.workdir = os.path.abspath(opts.workdir)

    setup_logging(opts)

    if not opts.workdir:
        tempfile.tempdir = opts.tmp

        # create the temporary directory for lorax
        tempdir = tempfile.mkdtemp(prefix="lorax.", dir=tempfile.gettempdir())
    else:
        tempdir = opts.workdir
        if not os.path.exists(tempdir):
            os.makedirs(tempdir)

    installtree = os.path.join(tempdir, "installtree")
    if not os.path.exists(installtree):
        os.mkdir(installtree)
    dnftempdir = os.path.join(tempdir, "dnf")
    if not os.path.exists(dnftempdir):
        os.mkdir(dnftempdir)

    dnfbase = get_dnf_base_object(installtree, opts.source, opts.mirrorlist,
                                  dnftempdir, opts.proxy, opts.version, opts.cachedir)

    if dnfbase is None:
        print("error: unable to create the dnf base object", file=sys.stderr)
        if not opts.workdir:
            shutil.rmtree(tempdir)
        sys.exit(1)

    parsed_add_template_vars = {}
    for kv in opts.add_template_vars:
        k, t, v = kv.partition('=')
        if t == '':
            raise ValueError("Missing '=' for key=value in " % kv)
        parsed_add_template_vars[k] = v

    parsed_add_arch_template_vars = {}
    for kv in opts.add_arch_template_vars:
        k, t, v = kv.partition('=')
        if t == '':
            raise ValueError("Missing '=' for key=value in " % kv)
        parsed_add_arch_template_vars[k] = v

    # run lorax
    lorax = pylorax.Lorax()
    lorax.configure(conf_file=opts.config)
    lorax.conf.set("lorax", "logdir", os.path.dirname(opts.logfile))
    lorax.run(dnfbase, opts.product, opts.version, opts.release,
              opts.variant, opts.bugurl, opts.isfinal,
              workdir=tempdir, outputdir=outputdir, buildarch=opts.buildarch,
              volid=opts.volid, domacboot=opts.domacboot, doupgrade=opts.doupgrade,
              installpkgs=opts.installpkgs,
              add_templates=opts.add_templates,
              add_template_vars=parsed_add_template_vars,
              add_arch_templates=opts.add_arch_templates,
              add_arch_template_vars=parsed_add_arch_template_vars,
              remove_temp=True)


def get_dnf_base_object(installroot, repositories, mirrorlists=None,
                        tempdir="/var/tmp", proxy=None, releasever="21",
                        cachedir=None):
    """ Create a dnf Base object and setup the repositories and installroot

        :param string installroot: Full path to the installroot
        :param list repositories: List of repositories to use for the installation
        :param list mirrorlist: List of mirrors to use
        :param string tempdir: Path of temporary directory
        :param string proxy: http proxy to use when fetching packages
        :param string releasever: Release version to pass to dnf
        :param string cachedir: Directory to use for caching packages

        If tempdir is not set /var/tmp is used.
        If cachedir is None a dnf.cache directory is created inside tmpdir
    """
    def sanitize_repo(repo):
        """Convert bare paths to file:/// URIs, and silently reject protocols unhandled by yum"""
        if repo.startswith("/"):
            return "file://{0}".format(repo)
        elif any(repo.startswith(p) for p in ('http://', 'https://', 'ftp://', 'file://')):
            return repo
        else:
            return None

    mirrorlists = mirrorlists or []

    # sanitize the repositories
    repositories = map(sanitize_repo, repositories)
    mirrorlists = map(sanitize_repo, mirrorlists)

    # remove invalid repositories
    repositories = filter(bool, repositories)
    mirrorlists = filter(bool, mirrorlists)

    if not cachedir:
        cachedir = os.path.join(tempdir, "dnf.cache")
    if not os.path.isdir(cachedir):
        os.mkdir(cachedir)

    logdir = os.path.join(tempdir, "dnf.logs")
    if not os.path.isdir(logdir):
        os.mkdir(logdir)

    dnfbase = dnf.Base()
    conf = dnfbase.conf
    conf.logdir = logdir
    conf.cachedir = cachedir

    # Turn off logging to the console
    conf.debuglevel = 10
    conf.errorlevel = 0
    dnfbase.logging.setup_from_dnf_conf(conf)

    conf.releasever = releasever
    conf.installroot = installroot
    conf.prepend_installroot('persistdir')
    conf.tsflags.append('nodocs')

    if proxy:
        conf.proxy = proxy

    # add the repositories
    for i, r in enumerate(repositories):
        if "SRPM" in r or "srpm" in r:
            log.info("Skipping source repo: %s", r)
            continue
        repo_name = "lorax-repo-%d" % i
        repo = dnf.repo.Repo(repo_name, cachedir)
        repo.baseurl = [r]
        if proxy:
            repo.proxy = proxy
        repo.enable()
        dnfbase.repos.add(repo)
        log.info("Added '%s': %s", repo_name, r)
        log.info("Fetching metadata...")
        try:
            repo.load()
        except dnf.exceptions.RepoError as e:
            log.error("Error fetching metadata for %s: %s", repo_name, e)
            return None

    # add the mirrorlists
    for i, r in enumerate(mirrorlists):
        if "SRPM" in r or "srpm" in r:
            log.info("Skipping source repo: %s", r)
            continue
        repo_name = "lorax-mirrorlist-%d" % i
        repo = dnf.repo.Repo(repo_name, cachedir)
        repo.mirrorlist = r
        if proxy:
            repo.proxy = proxy
        repo.enable()
        dnfbase.repos.add(repo)
        log.info("Added '%s': %s", repo_name, r)
        log.info("Fetching metadata...")
        try:
            repo.load()
        except dnf.exceptions.RepoError as e:
            log.error("Error fetching metadata for %s: %s", repo_name, e)
            return None

    dnfbase.fill_sack(load_system_repo=False)
    dnfbase.read_comps()

    return dnfbase


if __name__ == "__main__":
    main(sys.argv)