#!/usr/bin/python -tt
#
# livecd-creator : Creates Live CD based for Fedora.
#
# Copyright 2007, 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; version 2 of the License.
#
# 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 Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import os
import os.path
import glob
import shutil
import stat
import subprocess
import sys
import time
import optparse
import logging

import imgcreate
from imgcreate.fs import makedirs

class Usage(Exception):
    def __init__(self, msg = None, no_error = False):
        Exception.__init__(self, msg, no_error)

class LiveEFIImageCreator(imgcreate.LiveImageCreator):

    def _get_mkisofs_options(self, isodir):
        options = [ "-b", "isolinux/isolinux.bin",
                    "-c", "isolinux/boot.cat",
                    "-no-emul-boot", "-boot-info-table",
                    "-boot-load-size", "4" ]
        if os.path.exists(isodir + "/isolinux/efiboot.img"):
            options.extend([ "-eltorito-alt-boot",
                             "-e", "isolinux/efiboot.img",
                             "-no-emul-boot"])
        if os.path.exists(isodir + "/isolinux/macboot.img"):
            options.extend([ "-eltorito-alt-boot",
                             "-e", "isolinux/macboot.img",
                             "-no-emul-boot"])
        return options

    def __copy_efi_files(self, isodir):
        """ Copy the efi files into /EFI/BOOT/
            If any of them are missing, return False.
            requires:
              xen.efi
              gcdx64.efi
              vmlinuz
              initrd
        """
        fail = False
        missing = []
        # XXX: when adding multiple kernel support, vmlinuz and initrd needs to
        #      be suffixed with index
        files = [("/boot/efi/EFI/*/shim.efi", "/EFI/BOOT/BOOT%s.efi" % (self.efiarch,)),
                 ("/boot/efi/EFI/*/gcd*.efi", "/EFI/BOOT/grubx64.efi"),
                 ("/boot/efi/EFI/*/xen-*.efi", "/EFI/BOOT/xen.efi"),
                 ("/boot/efi/EFI/*/vmlinuz", "/EFI/BOOT/vmlinuz"),
                 ("/boot/efi/EFI/*/initrd-small.img", "/EFI/BOOT/initrd"),
                 ("/boot/efi/EFI/*/fonts/unicode.pf2", "/EFI/BOOT/fonts/"),
                ]
        makedirs(isodir+"/EFI/BOOT/fonts/")
        for src, dest in files:
            src_glob = glob.glob(self._instroot+src)
            if not src_glob:
                missing.append("Missing EFI file (%s)" % (src,))
                fail = True
            else:
                shutil.copy(src_glob[0], isodir+dest)
        map(logging.error, missing)
        return fail


    def __get_xen_efi_image_stanza(self, **args):
        if self._isDracut:
            args["rootlabel"] = "live:LABEL=%(fslabel)s" % args
        else:
            args["rootlabel"] = "CDLABEL=%(fslabel)s" % args
        return """[%(name)s%(index)s]
kernel=vmlinuz%(index)s root=%(rootlabel)s %(liveargs)s %(extra)s
ramdisk=initrd%(index)s

""" %args


    def __get_efi_image_stanza(self, **args):
        return """menuentry '%(long)s' --class qubes --class gnu-linux --class gnu --class os {
    chainloader /efi/boot/xen.efi placeholder %(name)s%(index)s
}
""" %args

    def __get_efi_image_stanzas(self, isodir, name):
        # FIXME: this only supports one kernel right now...

        kernel_options = self._get_kernel_options()
        checkisomd5 = self._has_checkisomd5()

        cfg = ""

        for index in range(0, 9):
            # only one supported anyway, so simply drop the suffix
            index = ""
            cfg += self.__get_efi_image_stanza(long = "Start " + self.product,
                                               index = index, name = "normal")
            if checkisomd5:
                cfg += self.__get_efi_image_stanza(
                                                   long = "Test this media & start " + self.product,
                                                   index = index, name = "check")
            cfg += """
submenu 'Troubleshooting -->' {
"""
            cfg += self.__get_efi_image_stanza(long = "Start " + self.product + " in basic graphics mode",
                                               index = index, name = "basicvideo")

            cfg+= """}
"""
            break

        return cfg

    def __get_xen_efi_image_stanzas(self, isodir, name):
        # FIXME: this only supports one kernel right now...

        kernel_options = self._get_kernel_options()
        checkisomd5 = self._has_checkisomd5()

        cfg = ""

        for index in range(0, 9):
            # only one supported anyway, so simply drop the suffix
            index = ""
            cfg += self.__get_xen_efi_image_stanza(fslabel = self.fslabel,
                                               liveargs = kernel_options,
                                               long = "Start " + self.product,
                                               extra = "", index = index,
                                               name = "normal")
            if checkisomd5:
                cfg += self.__get_xen_efi_image_stanza(fslabel = self.fslabel,
                                                   liveargs = kernel_options,
                                                   long = "Test this media & start " + self.product,
                                                   extra = "rd.live.check",
                                                   index = index, name = "check")
            cfg += self.__get_xen_efi_image_stanza(fslabel = self.fslabel,
                                               liveargs = kernel_options,
                                               long = "Start " + self.product + " in basic graphics mode",
                                               extra = "nomodeset", index = index,
                                               name = "basicvideo")

            break

        return cfg

    def __get_basic_xen_efi_config(self):
        return """
[global]
default=normal
"""

    def __get_basic_efi_config(self, **args):
        return """
set default="0"

function load_video {
  insmod efi_gop
  insmod efi_uga
  insmod video_bochs
  insmod video_cirrus
  insmod all_video
}

load_video
set gfxpayload=keep
insmod gzio
insmod part_gpt
insmod ext2

set timeout=%(timeout)d
### END /etc/grub.d/00_header ###

# do not use 'search' - root should be already set based on grub.efi location

### BEGIN /etc/grub.d/10_linux ###
""" %args

    def _configure_efi_bootloader(self, isodir):
        """Set up the configuration for an EFI bootloader"""
        if self.__copy_efi_files(isodir):
            shutil.rmtree(isodir + "/EFI")
            logging.warn("Failed to copy EFI files, no EFI Support will be included.")
            return

        cfg = self.__get_basic_efi_config(isolabel = self.fslabel,
                                          timeout = self._timeout)
        cfg += self.__get_efi_image_stanzas(isodir, self.name)

        xen_cfg = self.__get_basic_xen_efi_config()
        xen_cfg += self.__get_xen_efi_image_stanzas(isodir, self.name)

        cfgf = open(isodir + "/EFI/BOOT/grub.cfg", "w")
        cfgf.write(cfg)
        cfgf.close()

        xen_cfgf = open(isodir + "/EFI/BOOT/xen.cfg", "w")
        xen_cfgf.write(xen_cfg)
        xen_cfgf.close()

    def _generate_efiboot(self, isodir):
        """Generate EFI boot images."""
        if not glob.glob(self._instroot+"/boot/efi/EFI/*/xen-*.efi"):
            logging.error("Missing xen-*.efi, skipping efiboot.img creation.")
            return

        subprocess.call(["mkefiboot", "--label", "QUBESEFI", isodir + "/EFI/BOOT",
                         isodir + "/isolinux/efiboot.img"])
        # FIXME: replace icon
        # FIXME: this is broken for many reasons:
        #        - mkefiboot generates unnecessary big image (about 4 times bigger
        #          than required) - the bug is in mkmacboot function:
        #          size = estimate_size(bootdir, graft=graft) * 2
        #                               ^^^^^^^^^^^^^^^^^^^^ already counted twice
        #        - mkefiboot -a assumes that the loader is grub.efi
        #        - it isn't clear whether xen.efi would even work on Apple
        subprocess.call(["mkefiboot", "-a", isodir + "/EFI/BOOT",
                         isodir + "/isolinux/macboot.img", "-l", self.product,
                         "-n", "/usr/share/pixmaps/bootloader/fedora-media.vol",
                         "-i", "/usr/share/pixmaps/bootloader/fedora.icns",
                         "-p", self.product])


def parse_options(args):
    parser = optparse.OptionParser()

    imgopt = optparse.OptionGroup(parser, "Image options",
                                  "These options define the created image.")
    imgopt.add_option("-c", "--config", type="string", dest="kscfg",
                      help="Path or url to kickstart config file")
    imgopt.add_option("-b", "--base-on", type="string", dest="base_on",
                      help="Add packages to an existing live CD iso9660 image.")
    imgopt.add_option("-f", "--fslabel", type="string", dest="fslabel",
                      help="File system label (default based on config name)")
    imgopt.add_option("", "--title", type="string", dest="title",
                      help="Title used by syslinux.cfg file"),
    imgopt.add_option("", "--product", type="string", dest="product",
                      help="Product name used in syslinux.cfg boot stanzas and countdown"),
    # Provided for img-create compatibility
    imgopt.add_option("-n", "--name", type="string", dest="fslabel",
                      help=optparse.SUPPRESS_HELP)
    imgopt.add_option("-p", "--plugins", action="store_true", dest="plugins",
                      help="Use yum plugins during image creation",
                      default=False)
    imgopt.add_option("", "--image-type", type="string", dest="image_type",
                      help=optparse.SUPPRESS_HELP)
    imgopt.add_option("", "--compression-type", type="string", dest="compress_type",
                      help="Compression type recognized by mksquashfs "
                           "(default xz needs a 2.6.38+ kernel, gzip works "
                           "with all kernels, lzo needs a 2.6.36+ kernel, lzma "
                           "needs custom kernel.) Set to 'None' to force read "
                           "from base_on.",
                      default="xz")
    imgopt.add_option("", "--releasever", type="string", dest="releasever",
                      default=None,
                      help="Value to substitute for $releasever in kickstart repo urls")
    parser.add_option_group(imgopt)

    # options related to the config of your system
    sysopt = optparse.OptionGroup(parser, "System directory options",
                                  "These options define directories used on your system for creating the live image")
    sysopt.add_option("-t", "--tmpdir", type="string",
                      dest="tmpdir", default="/var/tmp",
                      help="Temporary directory to use (default: /var/tmp)")
    sysopt.add_option("", "--cache", type="string",
                      dest="cachedir", default=None,
                      help="Cache directory to use (default: private cache")
    sysopt.add_option("", "--cacheonly", action="store_true",
                      dest="cacheonly", default=False,
                      help="Work offline from cache, use together with --cache (default: False)")
    sysopt.add_option("", "--nocleanup", action="store_true",
                      dest="nocleanup", default=False,
                      help="Skip cleanup of temporary files")

    parser.add_option_group(sysopt)

    imgcreate.setup_logging(parser)

    # debug options not recommended for "production" images
    # Start a shell in the chroot for post-configuration.
    parser.add_option("-l", "--shell", action="store_true", dest="give_shell",
                      help=optparse.SUPPRESS_HELP)
    # Don't compress the image.
    parser.add_option("-s", "--skip-compression", action="store_true", dest="skip_compression",
                      help=optparse.SUPPRESS_HELP)
    parser.add_option("", "--skip-minimize", action="store_true", dest="skip_minimize",
                      help=optparse.SUPPRESS_HELP)

    (options, args) = parser.parse_args()

    # Pretend to be a image-creator if called with that name
    if not options.image_type:
        if sys.argv[0].endswith('image-creator'):
            options.image_type = 'image'
        else:
            options.image_type = 'livecd'
    if options.image_type not in ('livecd', 'image'):
        raise Usage("'%s' is not a recognized image type" % options.image_type)

    # image-create compatibility: Last argument is kickstart file
    if len(args) == 1:
        options.kscfg = args.pop()
    if len(args):
        raise Usage("Extra arguments given")

    if not options.kscfg or not os.path.exists(options.kscfg):
        raise Usage("Kickstart file must be provided")
    if options.base_on and not os.path.isfile(options.base_on):
        raise Usage("Image file '%s' does not exist" %(options.base_on,))
    if options.image_type == 'livecd':
        if options.fslabel and len(options.fslabel) > imgcreate.FSLABEL_MAXLEN:
            raise Usage("CD labels are limited to 32 characters")
        if options.fslabel and options.fslabel.find(" ") != -1:
            raise Usage("CD labels cannot contain spaces.")

    return options

def main():
    try:
        options = parse_options(sys.argv[1:])
    except Usage, (msg, no_error):
        if no_error:
            out = sys.stdout
            ret = 0
        else:
            out = sys.stderr
            ret = 2
        if msg:
            print >> out, msg
        return ret

    if os.geteuid () != 0:
        print >> sys.stderr, "You must run %s as root" % sys.argv[0]
        return 1

    if options.fslabel:
        fslabel = options.fslabel
        name = fslabel
    else:
        name = imgcreate.build_name(options.kscfg, options.image_type + "-")

        fslabel = imgcreate.build_name(options.kscfg,
                                        options.image_type + "-",
                                        maxlen = imgcreate.FSLABEL_MAXLEN,
                                        suffix = "%s-%s" %(os.uname()[4], time.strftime("%Y%m%d%H%M")))

        logging.info("Using label '%s' and name '%s'" % (fslabel, name))

    if options.title:
        title = options.title
    else:
        try:
            title = " ".join(name.split("-")[:2])
            title = title.title()
        except:
            title = "Linux"
    if options.product:
        product = options.product
    else:
        try:
            product = " ".join(name.split("-")[:2])
            product = product.title()
        except:
            product = "Linux"
    logging.info("Using title '%s' and product '%s'" % (title, product))

    ks = imgcreate.read_kickstart(options.kscfg)
    if not ks.handler.repo.seen:
        print >> sys.stderr, "Kickstart (%s) must have at least one repository." % (options.kscfg)
        return 1

    try:
        if options.image_type == 'livecd':
            creator = LiveEFIImageCreator(ks, name,
                                            fslabel=fslabel,
                                            releasever=options.releasever,
                                            tmpdir=os.path.abspath(options.tmpdir),
                                            useplugins=options.plugins,
                                            title=title, product=product,
                                            cacheonly=options.cacheonly,
                                            docleanup=not options.nocleanup)
        elif options.image_type == 'image':
            creator = imgcreate.LoopImageCreator(ks, name,
                                            fslabel=fslabel,
                                            releasever=options.releasever,
                                            useplugins=options.plugins,
                                            tmpdir=os.path.abspath(options.tmpdir),
                                            cacheonly=options.cacheonly,
                                            docleanup=not options.nocleanup)
    except imgcreate.CreatorError as e:
        logging.error(u"%s creation failed: %s", options.image_type, e)
        return 1

    creator.compress_type = options.compress_type
    creator.skip_compression = options.skip_compression
    creator.skip_minimize = options.skip_minimize
    if options.cachedir:
        options.cachedir = os.path.abspath(options.cachedir)

    try:
        creator.mount(options.base_on, options.cachedir)

        # fix /dev
        os.mknod(os.path.join(
            creator._instroot, 'dev/loop-control'), 0666 | stat.S_IFBLK, os.makedev(10, 237))
        for i in range(8):
            os.mknod(os.path.join(creator._instroot, 'dev/loop{}'.format(i)),
                0666 | stat.S_IFBLK, os.makedev(7, i))

        creator.install()
        creator.configure()
        if options.give_shell:
            print "Launching shell. Exit to continue."
            print "----------------------------------"
            creator.launch_shell()
        creator.unmount()
        creator.package()
    except imgcreate.CreatorError, e:
        logging.error(u"Error creating Live CD : %s" % e)
        return 1
    finally:
        creator.cleanup()

    return 0

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