#!/usr/bin/python

#
# lorax
#
# Copyright (C) 2009  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")
yum_log = logging.getLogger("yum")


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

import yum
# This is a bit of a hack to short circuit yum's internal logging
# handler setup.  We already set one up so we don't need it to run.
yum.logginglevels._added_handlers = True
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)

    # yum logging
    yum_log.setLevel(logging.DEBUG)
    logfile = os.path.abspath(os.path.dirname(opts.logfile))+"/yum.log"
    fh = logging.FileHandler(filename=logfile, mode="w")
    fh.setLevel(logging.DEBUG)
    yum_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("-e", "--excludepkgs", default=[],
            action="append", metavar="STRING",
            help="package glob to exclude (may be listed multiple times)")
    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", dest="domacboot")
    optional.add_option("--nomacboot", help="",
            action="store_false", default=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("--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=[])
    optional.add_option("--noverifyssl", action="store_true", default=False,
                        help="Do not verify SSL certificates")

    # 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 os.path.exists(outputdir):
        parser.error("output directory should not exist.")

    opts.logfile = os.path.abspath(opts.logfile)

    if opts.config and not os.path.exists(opts.config):
        parser.error("config file %s doesn't exist." % opts.config)

    setup_logging(opts)

    tempfile.tempdir = opts.tmp

    # create the temporary directory for lorax
    tempdir = tempfile.mkdtemp(prefix="lorax.", dir=tempfile.gettempdir())

    # create the yumbase object
    installtree = os.path.join(tempdir, "installtree")
    os.mkdir(installtree)
    yumtempdir = os.path.join(tempdir, "yum")
    os.mkdir(yumtempdir)

    yb = get_yum_base_object(installtree, opts.source, opts.mirrorlist,
                             yumtempdir, opts.proxy, opts.excludepkgs,
                             not opts.noverifyssl)

    if yb is None:
        print("error: unable to create the yumbase object", file=sys.stderr)
        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.run(yb, 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_yum_base_object(installroot, repositories, mirrorlists=[],
                        tempdir="/var/tmp", proxy=None, excludepkgs=[],
                        sslverify=True):

    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

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

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

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

    yumconf = os.path.join(tempdir, "yum.conf")
    c = ConfigParser.ConfigParser()

    # add the main section
    section = "main"
    data = {"cachedir": cachedir,
            "keepcache": 0,
            "gpgcheck": 0,
            "plugins": 0,
            "reposdir": "",
            "tsflags": "nodocs"}

    if proxy:
        data["proxy"] = proxy

    if sslverify == False:
        data["sslverify"] = "0"

    if excludepkgs:
        data["exclude"] = " ".join(excludepkgs)

    c.add_section(section)
    map(lambda (key, value): c.set(section, key, value), data.items())

    # add the main repository - the first repository from list
    section = "lorax-repo"
    data = {"name": "lorax repo",
            "baseurl": repositories[0],
            "enabled": 1}

    c.add_section(section)
    map(lambda (key, value): c.set(section, key, value), data.items())

    # add the extra repositories
    for n, extra in enumerate(repositories[1:], start=1):
        section = "lorax-extra-repo-{0:d}".format(n)
        data = {"name": "lorax extra repo {0:d}".format(n),
                "baseurl": extra,
                "enabled": 1}

        c.add_section(section)
        map(lambda (key, value): c.set(section, key, value), data.items())

    # add the mirrorlists
    for n, mirror in enumerate(mirrorlists, start=1):
        section = "lorax-mirrorlist-{0:d}".format(n)
        data = {"name": "lorax mirrorlist {0:d}".format(n),
                "mirrorlist": mirror,
                "enabled": 1 }

        c.add_section(section)
        map(lambda (key, value): c.set(section, key, value), data.items())

    # write the yum configuration file
    with open(yumconf, "w") as f:
        c.write(f)

    # create the yum base object
    yb = yum.YumBase()

    yb.preconf.fn = yumconf
    yb.preconf.root = installroot
    #yb.repos.setCacheDir(cachedir)

    # Turn on as much yum logging as we can
    yb.preconf.debuglevel = 6
    yb.preconf.errorlevel = 6
    yb.logger.setLevel(logging.DEBUG)
    yb.verbose_logger.setLevel(logging.DEBUG)

    return yb


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