#!/usr/bin/python -ttd
# version date: 2011-01-28 02:41:19 -0500
#
# edit-liveos: Edit a LiveOS to insert files or to clone an instance onto a new
#                   iso image file.
#
# Copyright 2009, Red Hat  Inc.
# Written by Perry Myers <pmyers at redhat.com> & David Huff <dhuff at redhat.com>
#   Cloning code added by Frederick Grose <fgrose at sugarlabs.org>
#
#
# 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 sys
import stat
import tempfile
import shutil
import subprocess
import optparse
import logging

from imgcreate.debug import *
from imgcreate.fs import *
from imgcreate.errors import *
from imgcreate.live import *
from imgcreate.creator import *

class ExistingSparseLoopbackDisk(SparseLoopbackDisk):
    """don't want to expand the disk"""
    def __init__(self, lofile, size):
        SparseLoopbackDisk.__init__(self, lofile, size)

    def create(self):
        #self.expand(create = True)
        LoopbackDisk.create(self)

class LiveImageEditor(LiveImageCreator):
    """class for editing LiveOS images.

    We need an instance of LiveImageCreator, however, we do not have a kickstart
    file and we may not need to create a new image.  We just want to reuse some
    of LiveImageCreators methods on an existing LiveOS image.

    """

    def __init__(self, name):
        """Initialize a LiveImageEditor instance.

        creates a dummy instance of LiveImageCreator
        We do not initialize any sub classes b/c we have no ks file.

        """
        self.name = name

        self.tmpdir = "/var/tmp"
        """The directory in which all temporary files will be created."""

        self.clone = False
        """Signals when to copy a running LiveOS image as base."""

        self._include = None
        """A string of file or directory paths to include in __copy_img_root."""

        self._builder = "someone"
        """The name of the Remix builder for _branding."""

        self.compress_type = "gzip"
        """mksquashfs compressor to use."""

        self.skip_compression = False
        """Controls whether to use squashfs to compress the image."""

        self.skip_minimize = False
        """Controls whether an image minimizing snapshot should be created."""

        self._isofstype = "iso9660"
        self.__isodir = None

        self._ImageCreator__builddir = None
        """working directory"""

#        self._ImageCreator_instroot = None
        """where the extfs.img is mounted for modification"""

        self._ImageCreator_outdir = None
        """where final iso gets written"""

        self._ImageCreator__bindmounts = []

        # variables are set with the next assignment????
#        self._LoopImageCreator__imagedir = None
        """dir for the extfs.img"""

        self._LoopImageCreator__blocksize = 4096
        self._LoopImageCreator__fslabel = None
        self._LoopImageCreator__instloop = None
        self._LoopImageCreator__fstype = None
        self._LoopImageCreator__image_size = None

        self.__image_size = 4096L * 1024 * 1024
        self.__fstype = "ext4"
        self.__blocksize = 4096
        self.__instroot = None

        self._LiveImageCreatorBase__isodir = None
        """directory where the iso is staged"""

    # properties
    def __get_image(self):
        if self._LoopImageCreator__imagedir is None:
            self.__ensure_builddir()
            self._LoopImageCreator__imagedir = \
                tempfile.mkdtemp(dir = os.path.abspath(self.tmpdir),
                                                       prefix = self.name + "-")
        rtn = self._LoopImageCreator__imagedir + "/ext3fs.img"
        return rtn
    _image = property(__get_image)
    """The location of the image file or filesystem root."""

    def _get_fstype(self):
        if self.clone:
            # FIXME detect the fstype
            self._LoopImageCreator__fstype = "ext4"
        else:
            dev_null = os.open("/dev/null", os.O_WRONLY)
            try:
                out = subprocess.Popen(["/sbin/blkid", self._image],
                                       stdout = subprocess.PIPE,
                                       stderr = dev_null).communicate()[0]
                for word in out.split():
                    if word.startswith("TYPE"):
                        self._LoopImageCreator__fstype = word.split("=")[1].strip("\"")

            except IOError, e:
                raise CreatorError("Failed to determine fsimage TYPE: %s" % e )

    # _fstype = property(__get_fstype, __set_fstype)

    def _get_fslabel(self):
        if self.clone:
            # FIXME detect the fslabel
            self._LoopImageCreator__fslabel = self.name
        else:
            dev_null = os.open("/dev/null", os.O_WRONLY)
            try:
                out = subprocess.Popen(["/sbin/e2label", self._image],
                                       stdout = subprocess.PIPE,
                                       stderr = dev_null).communicate()[0]

                self._LoopImageCreator__fslabel = out.strip()

            except IOError, e:
                raise CreatorError("Failed to determine fsimage LABEL: %s" % e )

    def __ensure_builddir(self):
        if not self._ImageCreator__builddir is None:
            return

        try:
            self._ImageCreator__builddir = tempfile.mkdtemp(dir =  os.path.abspath(self.tmpdir),
                                               prefix = "edit-liveos-")
        except OSError, (err, msg):
            raise CreatorError("Failed create build directory in %s: %s" %
                               (self.tmpdir, msg))

    def _run_script(self, script):

        (fd, path) = tempfile.mkstemp(prefix = "script-",
                                          dir = self._instroot + "/tmp")

        logging.debug("copying script to install root: %s" % path)
        shutil.copy(os.path.abspath(script), path)
        os.close(fd)
        os.chmod(path, 0700)

        script = "/tmp/" + os.path.basename(path)

        try:
            subprocess.call([script], preexec_fn = self._chroot)
        except OSError, e:
            raise CreatorError("Failed to execute script %s, %s " % (script, e))
        finally:
            os.unlink(path)

    def mount(self, base_on, cachedir = None):
        """mount existing file system.

        We have to override mount b/c we many not be creating an new install
        root nor do we need to setup the file system, i.e., makedirs(/etc/,
        /boot, ...), nor do we want to overwrite fstab, or create selinuxfs.

        We also need to get some info about the image before we can mount it.

        base_on --  the <LIVEIMG.src> a LiveOS.iso file or an attached LiveOS
                    device, such as, /dev/live for a currently running image.

        cachedir -- a directory in which to store a Yum cache;
                    Not used in edit-liveos.

        """

        if not base_on:
            raise CreatorError("No base LiveOS image specified.")

        self.__ensure_builddir()

        self._ImageCreator_instroot = self._ImageCreator__builddir + "/install_root"
        self._LoopImageCreator__imagedir = self._ImageCreator__builddir + "/ex"
        self._ImageCreator_outdir = self._ImageCreator__builddir + "/out"

        makedirs(self._ImageCreator_instroot)
        makedirs(self._LoopImageCreator__imagedir)
        makedirs(self._ImageCreator_outdir)

        if self.clone:
            # Need to clone base_on into ext3fs.img at this point
            self._base_on(base_on)
        else:
            LiveImageCreator._base_on(self, base_on)

        self._LoopImageCreator__image_size = os.stat(self._image)[stat.ST_SIZE]

        self._get_fstype()
        self._get_fslabel()
        self.fslabel = self._LoopImageCreator__fslabel

        self._LoopImageCreator__instloop = \
            ExtDiskMount(ExistingSparseLoopbackDisk(self._image,
                                        self._LoopImageCreator__image_size),
                         self._ImageCreator_instroot,
                         self._fstype,
                         self._LoopImageCreator__blocksize,
                         self.fslabel)
        try:
            self._LoopImageCreator__instloop.mount()
        except MountError, e:
            raise CreatorError("Failed to loopback mount '%s' : %s" %
                               (self._image, e))

        cachesrc = cachedir or (self._ImageCreator__builddir + "/yum-cache")
        makedirs(cachesrc)

        for (f, dest) in [("/sys", None), ("/proc", None),
                          ("/dev/pts", None), ("/dev/shm", None),
                          (cachesrc, "/var/cache/yum")]:
            self._ImageCreator__bindmounts.append(BindChrootMount(f, self._instroot, dest))

        self._do_bindmounts()

        os.symlink("../proc/mounts", self._instroot + "/etc/mtab")

        self.__copy_img_root(base_on)
        self._brand(self._builder)

    def _base_on(self, base_on):
        """Clone the running LiveOS image as the basis for the new image."""

        self.__instloop = ExtDiskMount(SparseLoopbackDisk(self._image,
                                                          self.__image_size),
                                       self._instroot,
                                       self.__fstype,
                                       self.__blocksize,
                                       self.fslabel)
        try:
            self.__instloop.mount()
        except MountError, e:
            raise CreatorError("Failed to loopback mount '%s' : %s" %
                               (self._image, e))

        subprocess.call(['rsync', '-ptgorlHASx', '--specials', '--progress',
                         '--include', '/*/',
                         '--include', '/etc/blkid/',
                         '--exclude', '*',
                         '/', self._instroot])
        subprocess.call(['sync'])
        subprocess.call(['rsync', '-ptgorlHASx', '--specials', '--progress',
                         '--exclude', '/etc/mtab',
                         '--exclude', '/etc/blkid*',
                         '--exclude', '/dev*',
                         '--exclude', '/proc*',
                         '--exclude', '/home*',
                         '--exclude', '/media*',
                         '--exclude', '/mnt*',
                         '--exclude', '/sys*',
                         '--exclude', '/tmp*',
                         '--exclude', '/.liveimg*',
                         '--exclude', '/.autofsck',
                         '/', self._instroot])
        subprocess.call(['sync'])

        self._ImageCreator__create_minimal_dev()

        self.__instloop.cleanup()

    def _mount_instroot(self, base_on = None):
        self.__imgdir = self._mkdtemp()

        self.__instloop = ExtDiskMount(self._image,
                                       self._instroot,
                                       self.__fstype,
                                       self.__blocksize,
                                       self.fslabel)

        try:
            self.__instloop.mount()
        except MountError, e:
            raise CreatorError("Failed to mount '%s' : %s" %
                               (self._image, e))

    def __copy_img_root(self, base_on):
        """helper function to copy root content of the base LiveIMG to ISOdir"""

        ignore_list = ['ext3fs.img', 'squashfs.img', 'osmin.img', 'home.img',
                       'overlay-*']

        if self.clone:
            ignore_list.remove('home.img')
            includes = 'boot, /EFI, /syslinux, /LiveOS'
            if self._include:
                includes += ", " + self._include
            imgmnt = DiskMount(RawDisk(0, base_on), self._mkdtemp())
        else:
            imgmnt = DiskMount(LoopbackDisk(base_on, 0), self._mkdtemp())

        self._LiveImageCreatorBase__isodir = self._ImageCreator__builddir + "/iso"

        try:
            imgmnt.mount()

            # include specified files or directories
            if self.clone:
                dst = self._LiveImageCreatorBase__isodir
                print includes
                for fd in includes.split(', /'):
                    src = os.path.join(imgmnt.mountdir, fd)
                    if os.path.isfile(src):
                        shutil.copy2(src, os.path.join(dst, fd))
                    elif os.path.isdir(src):
                        shutil.copytree(src, os.path.join(dst, fd), symlinks=True,
                                        ignore=shutil.ignore_patterns(*ignore_list))
            else:
                #copy over everything but squashfs.img or ext3fs.img
                shutil.copytree(imgmnt.mountdir,
                                self._LiveImageCreatorBase__isodir,
                                ignore=shutil.ignore_patterns(*ignore_list))
        except MountError, e:
            raise CreatorError("Failed to mount '%s' : %s" %
                               (base_on, e))

        finally:
            subprocess.call(['sync'])
            imgmnt.cleanup()

    def _brand (self, _builder):
        """Adjust the image branding to show its variation from original
        source by builder and build date."""

        dt = time.strftime("%d-%b-%Y")

        # Get build name from boot configuration file.
        try:
            cfgf = open(self._LiveImageCreatorBase__isodir +
                        "/syslinux/syslinux.cfg", "r")
        except IOError:
            cfgf = open(self._LiveImageCreatorBase__isodir +
                        "/syslinux/extlinux.conf", "r")
        for line in cfgf:
            i = line.find("Welcome to ")
            if i > -1:
                self.name = line[i+11:-2]
                break
        cfgf.close()

        # Update fedora-release message with Remix details.
        releasefiles = '/etc/fedora-release, /etc/generic-release'
        if self._releasefile:
            releasefiles += ', ' + self._releasefile
        for f in releasefiles.split(', '):
            if os.path.exists(f):
                try:
                    text = open(self._instroot + f, "r").read()
                except IOError:
                    raise CreatorError("Failed to open '%s' : %s" % (cfgf, e))

                text += "  Remix of " + self.name + " by " + _builder + \
                        " on " + dt
                open(self._instroot + f, "w").write(text)
        self.name += " " + dt + " Remix"
        self.fslabel += '-' + os.uname()[4] + '-' + time.strftime("%Y%m%d.%H")

    def _configure_bootloader(self, isodir):
        """Restore the boot configuration files for an iso image boot."""

        bootfolder = os.path.join(isodir, 'isolinux')
        oldpath = os.path.join(isodir, 'syslinux')
        if os.path.exists(oldpath):
            os.rename(oldpath, bootfolder)

        cfgf = os.path.join(bootfolder, 'isolinux.cfg')
        for f in ['syslinux.cfg', 'extlinux.conf']:
            src = os.path.join(bootfolder, f)
            if os.path.exists(src):
                os.rename(src, cfgf)

        args = ['/bin/sed', '-i',
                '-e', 's/Welcome to .*/Welcome to ' + self.name + '!/',
                '-e', 's/root=[^ ]*/root=live:CDLABEL=' + self.fslabel + '/',
                '-e', 's/rootfstype=[^ ]* [^ ]*/rootfstype=auto ro/',
                '-e', 's/liveimg .* quiet/liveimg quiet/', cfgf]

        dev_null = os.open("/dev/null", os.O_WRONLY)
        try:
            subprocess.Popen(args,
                             stdout = subprocess.PIPE,
                             stderr = dev_null).communicate()[0]
            return 0

        except IOError, e:
            raise CreatorError("Failed to configure the bootloader file: %s" % e)
            return 1

def parse_options(args):
    parser = optparse.OptionParser(usage = "%prog [-n=<name>] \
                           \n\r                   [-o=<output>] \
                           \n\r                   [-s=<script.sh>] \
                           \n\r                   [-t=<tmpdir>] \
                           \n\r                   [-i=<includes>] \
                           \n\r                   [-r=<releasefile>] \
                           \n\r                   [--builder] \
                           \n\r                   [--clone] \
                           \n\r                   [--skip-compression] \
                           \n\r                   [--skip-minimize] \
                           \n\r                   <LIVEIMG.src>")

    parser.add_option("-n", "--name", type="string", dest="name",
                      help="name of new LiveOS (don't include .iso, it will be added)")

    parser.add_option("-o", "--output", type="string", dest="output",
                      help="specify directory for new iso file.")

    parser.add_option("-s", "--script", type="string", dest="script",
                      help="specify script to run chrooted in the LiveOS fsimage")

    parser.add_option("", "--clone", action="store_true", dest="clone",
                      help="Specify that source image is LiveOS block device.")

    parser.add_option("-i", "--include", type="string", dest="include",
                      help="Specify directory or file patterns to be included in copy_img_root.")

    parser.add_option("-r", "--releasefile", type="string", dest="releasefile",
                      help="Specify release file/s for branding.")

    parser.add_option("", "--builder", type="string", dest="builder",
                      help="Specify the builder of a Remix.")

    parser.add_option("-t", "--tmpdir", type="string",
                      dest="tmpdir", default="/var/tmp",
                      help="Temporary directory to use (default: /var/tmp)")

    parser.add_option("", "--skip-compression", action="store_true", dest="skip_compression")

    parser.add_option("", "--skip-minimize", action="store_true", dest="skip_minimize")

    setup_logging(parser)

    (options, args) = parser.parse_args()

    if len(args) != 1:
        parser.print_usage()
        sys.exit(1)

    print args[0]
    if stat.S_ISBLK(os.stat(args[0]).st_mode):
        options.clone = True

    return (args[0], options)

def rebuild_iso_symlinks(isodir):
    # remove duplicate files and rebuild symlinks to reduce iso size
    efi_vmlinuz = "%s/EFI/boot/vmlinuz0" % isodir
    isolinux_vmlinuz = "%s/isolinux/vmlinuz0" % isodir
    efi_initrd = "%s/EFI/boot/initrd0.img" % isodir
    isolinux_initrd = "%s/isolinux/initrd0.img" % isodir

    if os.path.exists(efi_vmlinuz):
        os.remove(efi_vmlinuz)
        os.remove(efi_initrd)
        os.symlink(isolinux_vmlinuz,efi_vmlinuz)
        os.symlink(isolinux_initrd,efi_initrd)

def main():
    # LiveOS set to <LIVEIMG.src>
    (LiveOS, options) = parse_options(sys.argv[1:])

    if os.geteuid () != 0:
        print >> sys.stderr, "You must run edit-liveos as root"
        return 1

    if options.name and options.name != os.path.basename(LiveOS):
        name = options.name
    else:
        name = os.path.basename(LiveOS) + ".edited"

    if options.output:
        output = options.output
    else:
        output = os.path.dirname(LiveOS)

    editor = LiveImageEditor(name)
    editor._include = options.include
    editor.clone = options.clone
    editor.tmpdir = options.tmpdir
    editor._builder = options.builder
    editor._releasefile = options.releasefile
    editor.skip_compression = options.skip_compression
    editor.skip_minimize = options.skip_minimize

    try:
        editor.mount(LiveOS, cachedir = None)
        if editor.clone:
            editor._configure_bootloader(editor._LiveImageCreatorBase__isodir)
            editor.name = editor.fslabel
#        if options.script:
#            print "Running edit script '%s'" % options.script
#            editor._run_script(options.script)
#        else:
#            print "Launching shell. Exit to continue."
#            print "----------------------------------"
#            editor.launch_shell()
        rebuild_iso_symlinks(editor._LiveImageCreatorBase__isodir)
        editor.unmount()
        if editor.clone:
            #FIXME fs.py.resize2fs defaults to /tmp (tmpfs on LiveOS)
            ## workaround for hardcoded /tmp in fs.py.resize2fs
            tmpdirhack = editor._mkdtemp()
            os.rmdir(tmpdirhack)
            try:
                shutil.copytree('/tmp', tmpdirhack)
            except:
                pass
            editor._ImageCreator__bindmounts = [BindChrootMount(
                                                   tmpdirhack,'', 'tmp')]
            editor._do_bindmounts()
        editor.package(output)
        editor.clone and editor._undo_bindmounts()
    except CreatorError, e:
        logging.error(u"Error editing LiveOS : %s" % e)
        return 1
    finally:
        editor.cleanup()

    return 0

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

#tracer = trace.Trace(trace=1)
#tracer.run('main()')
arch = rpmUtils.arch.getBaseArch()
if arch in ("i386", "x86_64"):
    LiveImageCreator = x86LiveImageCreator
elif arch in ("ppc",):
    LiveImageCreator = ppcLiveImageCreator
elif arch in ("ppc64",):
    LiveImageCreator = ppc64LiveImageCreator
else:
    raise CreatorError("Architecture not supported!")