#!/usr/bin/python -tt
#
# livecd-creator : Creates Live CD based for Fedora.
#
# 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 getopt
import os
import os.path
import sys
import tempfile
import time

class InstallationTarget:

    def base_on_iso(self, base_on):
        """helper function to extract ext3 file system from a live CD ISO"""
        
        # create directories for iso9660 and squashfs file systems from the live cd ISO
        try:
            os.mkdir(self.build_dir + "/base_on_iso")
            os.mkdir(self.build_dir + "/base_on_squashfs")
        except OSError:
            return False

        # setup loop device for iso9660 file system
        try:
            f = os.popen("/sbin/losetup -f")
            iso_loop_device = f.read().strip()
            f.close()
        except OSError:
            print "Cannot get loop device"
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False
        if os.system("/sbin/losetup %s %s"%(iso_loop_device, base_on)) != 0:
            print "Cannot setup %s for ISO to base on"%(self.loop_device)
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False

        # mount iso9660 file system
        if os.system("/bin/mount %s %s/base_on_iso"%(iso_loop_device, self.build_dir)) != 0:
            print "Cannot mount ISO to base on in %s/base_on_iso"%(self.build_dir)
            os.system("/sbin/losetup -d %s"%(iso_loop_device))
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False

        # setup loop device for squashfs file system
        try:
            f = os.popen("/sbin/losetup -f")
            squashfs_loop_device = f.read().strip()
            f.close()
        except OSError:
            print "Cannot get loop device"
            os.system("/bin/umount %s/base_on_iso"%(self.build_dir))
            os.system("/sbin/losetup -d %s"%(iso_loop_device))
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False        
        if os.system("/sbin/losetup %s %s/base_on_iso/squashfs.img"%(squashfs_loop_device, self.build_dir)) != 0:
            print "Cannot setup %s for squashfs on ISO to base on"%(self.loop_device)
            os.system("/bin/umount %s/base_on_iso"%(self.build_dir))
            os.system("/sbin/losetup -d %s"%(iso_loop_device))
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False

        # mount squashfs file system
        if os.system("/bin/mount -t squashfs -o ro %s %s/base_on_squashfs"%(squashfs_loop_device, self.build_dir)) != 0:
            print "Cannot mount squashfs from ISO to base on in %s/base_on_iso"%(self.build_dir)
            os.system("/sbin/losetup -d %s"%(squashfs_loop_device))
            os.system("/bin/umount %s/base_on_iso"%(self.build_dir))
            os.system("/sbin/losetup -d %s"%(iso_loop_device))
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False

        # copy the ext3 fs out
        if os.system("cp %s/base_on_squashfs/os.img %s/data/os.img"%(self.build_dir, self.build_dir)) != 0:
            print "Cannot copy os.img from squashfs from ISO to base on"
            os.system("/bin/umount %s/base_on_squashfs"%(self.build_dir))
            os.system("/sbin/losetup -d %s"%(squashfs_loop_device))
            os.system("/bin/umount %s/base_on_iso"%(self.build_dir))
            os.system("/sbin/losetup -d %s"%(iso_loop_device))
            try:
                os.rmdir(self.build_dir + "/base_on_squashfs")
                os.rmdir(self.build_dir + "/base_on_iso")
            except OSError:
                pass
            return False

        # unmount and tear down the mount points and loop devices used
        os.system("/bin/umount %s/base_on_squashfs"%(self.build_dir))
        os.system("/sbin/losetup -d %s"%(squashfs_loop_device))
        os.system("/bin/umount %s/base_on_iso"%(self.build_dir))
        os.system("/sbin/losetup -d %s"%(iso_loop_device))
        try:
            os.rmdir(self.build_dir + "/base_on_squashfs")
            os.rmdir(self.build_dir + "/base_on_iso")
        except OSError:
            pass
        
        return True

    def check_free_space(self, scratch_space_dir):
        """Check for space on defined build directory."""
        self.scratch_space_dir = scratch_space_dir
        varspace = os.statvfs(scratch_space_dir)

        """Based on the default block size, calculate free space, and hand it back to the caller."""
        calcfreevar = varspace[0]*varspace[4]
        return calcfreevar

    def setup(self, image_size, fs_label, base_on, scratch_space_dir):

        """setup target ext3 file system in preparation for an install"""

        # global variables needed
        self.fs_label = fs_label
        self.have_unmounted = False
        self.loop_device = ""

        # setup temporary build dirs
        appended_dir = scratch_space_dir + "build-" + os.path.basename(tempfile.mktemp())
        self.build_dir = appended_dir
        try:
            os.mkdir(self.build_dir)
        except OSError:
            print "Cannot create build directory at %s"
            return False

        try:
            os.mkdir(self.build_dir + "/out")
            os.mkdir(self.build_dir + "/out/boot")
            os.mkdir(self.build_dir + "/out/boot/grub")
            os.mkdir(self.build_dir + "/out/sysroot")
            os.mkdir(self.build_dir + "/data")
            os.mkdir(self.build_dir + "/data/sysroot")
            os.mkdir(self.build_dir + "/install_root")
            os.mkdir(self.build_dir + "/yum-cache")
        except OSError:
            print "Cannot create directory"
            self.teardown()
            return False

        if base_on != "":
            # get backing ext3 image if we're based this build on an existing live CD ISO
            if not self.base_on_iso(base_on):
                self.teardown()
                return False
        else:
            # otherwise create the image
            if os.system("dd if=/dev/zero of=%s/data/os.img bs=512 count=1 seek=%d 2> /dev/null"%(self.build_dir, image_size/512)) != 0:
                print "Cannot create sparse file of %d MiB for OS image"%(image_size/1024/1024)
                self.teardown()
                return False

        # get a loop device, self.loop_device, to point to the ext3 image we're using
        try:
            f = os.popen("/sbin/losetup -f")
            self.loop_device = f.read().strip()
            f.close()
        except OSError:
            print "Cannot get loop device"
            self.teardown()
            return False
        if os.system("/sbin/losetup %s %s/data/os.img"%(self.loop_device, self.build_dir)) != 0:
            print "Cannot setup %s for backing OS image"%(self.loop_device)
            self.teardown()
            return False

        # create the ext3 image unless we're reusing one from an existing live CD ISO
        if base_on == "":
            if os.system("/sbin/mkfs.ext3 -L \"%s\" -m 1 %s > /dev/null 2>&1"%(self.fs_label, self.loop_device)) != 0:
                print "Cannot create ext3 file system on OS image"
                self.teardown()
                return False

            if os.system("/sbin/tune2fs -c 0 -i 0 %s > /dev/null 2>&1"%(self.loop_device)) != 0:
                print "Cannot tune ext3 file system on OS image"
                self.teardown()
                return False

        # mount ext3 image at install_root/ so we can write to it
        if os.system("/bin/mount -t ext3 %s %s/install_root"%(self.loop_device, self.build_dir)) != 0:
            print "Cannot mount ext3 fs from OS image in %s/install_root"%(self.build_dir)
            self.teardown()
            return False

        if base_on == "":
            # create a few directories needed if it's a new image
            try:
                os.mkdir(self.build_dir + "/install_root/etc")
                os.mkdir(self.build_dir + "/install_root/boot")
                os.mkdir(self.build_dir + "/install_root/var")
                os.mkdir(self.build_dir + "/install_root/var/log")
                os.mkdir(self.build_dir + "/install_root/var/cache")
                os.mkdir(self.build_dir + "/install_root/var/cache/yum")
            except OSError:
                print "Cannot create directory"
                self.teardown()
                return False

        # bind mount system directories into install_root/
        for f in ["sys", "proc", "dev", "dev/pts", "selinux"]:
            try:
                os.makedirs(self.build_dir + "/install_root/%s"%f)
            except:
                pass
            if os.system("/bin/mount --bind /%s %s/install_root/%s"%(f, self.build_dir, f)) != 0:
                print "Cannot mount special file system %s"%f
                self.teardown()
                return False

        # make sure /etc/mtab is current inside install_root
        try:
            os.symlink("../proc/mounts", self.build_dir + "/install_root/etc/mtab")
        except OSError:
            print "Cannot create symlink %s/etc/mtab -> ../proc/mounts"%(self.build_dir)
            self.teardown()
            return False

        # write an /etc/fstab file
        fstab = open(self.build_dir + "/install_root/etc/fstab", "w")
        fstab.write("/dev/mapper/livecd-rw   /                       ext3    defaults,noatime 0 0\n")
        fstab.write("devpts                  /dev/pts                devpts  gid=5,mode=620  0 0\n")
        fstab.write("tmpfs                   /dev/shm                tmpfs   defaults        0 0\n")
        fstab.write("proc                    /proc                   proc    defaults        0 0\n")
        fstab.write("sysfs                   /sys                    sysfs   defaults        0 0\n")
        fstab.close()

        if os.system("/bin/mount --bind %s/yum-cache %s/install_root/var/cache/yum"%(self.build_dir, self.build_dir)) != 0:
            print "Cannot bind mount /var/cache/yum"
            self.teardown()
            return False

        # setup yum.conf (this file live _outside_ the intall root)
        yumconf = open(self.build_dir + "/yum.conf", "w")
        yumconf.write("[main]\n")
        yumconf.write("cachedir=/var/cache/yum\n")
        yumconf.write("debuglevel=1\n")
        yumconf.write("logfile=/var/log/yum.log\n")
        yumconf.write("pkgpolicy=newest\n")
        yumconf.write("distroverpkg=redhat-release\n")
        yumconf.write("tolerant=1\n")
        yumconf.write("exactarch=1\n")
        yumconf.write("retries=20\n")
        yumconf.write("obsoletes=1\n")
        yumconf.write("gpgcheck=0\n")
        yumconf.write("\n")
        yumconf.close()

        self.repo_num = 0
        return True


    def unmount(self):
        """detaches system bind mounts and install_root for the file system and tears down loop devices used"""
        
        try:
            os.unlink(self.build_dir + "/install_root/etc/mtab")
        except OSError:
            pass

        # TODO: need to handle when unmount fails because of lingering
        # processes that has open files
        if not self.have_unmounted:
            os.system("/bin/umount %s/install_root/var/cache/yum"%(self.build_dir))
            for g in ["selinux", "dev/pts", "dev", "proc", "sys"]:
                os.system("/bin/umount %s/install_root/%s"%(self.build_dir, g))
            if self.loop_device != "":
                os.system("/bin/umount %s"%(self.loop_device))
                os.system("/sbin/losetup -d %s"%(self.loop_device))

            try:
                os.rmdir(self.build_dir + "/install_root")
                os.unlink(self.build_dir + "/yum.conf")
            except OSError:
                pass
            os.system("/bin/rm -rf %s/yum-cache"%(self.build_dir))
            
            self.have_unmounted = True

    def teardown(self):
        """releases all resources and removes all temporary files"""

        # ensure we've detached the install_root
        self.unmount()

        # this removes all data used in the temporary build directory
        try:
            os.unlink(self.build_dir + "/data/os.img")
            os.rmdir(self.build_dir + "/data")
        except OSError:
            pass
        os.system("/bin/rm -rf %s/out"%(self.build_dir))
        #if "/var/tmp/livecd-creator/build-" in self.build_dir:
        if self.build_dir + "build-" in self.build_dir:
            os.system("/bin/rm -rf %s"%self.build_dir)

    def addRepository(self, name, url):
        """adds a yum repository to temporary yum.conf file used"""

        yumconf = open(self.build_dir + "/yum.conf", "a")
        yumconf.write("[lcdr_%s]\n"%name)
        yumconf.write("name=Live CD creator repo for '%s'\n"%name)
        yumconf.write("baseurl=%s\n"%url)
        yumconf.write("enabled=1\n")
        yumconf.write("gpgcheck=0\n")
        yumconf.write("\n")
        yumconf.close()

    def installPackages(self, packageList, excludePackageList):
        """install packages into target file system"""
        
        packages = ""
        for p in packageList:
            packages += " %s"%p

        os.system("/bin/touch %s/install_root/var/log/yum.log"%(self.build_dir))

        if os.system("/usr/bin/yum -y -c %s/yum.conf --disablerepo=* --enablerepo=lcdr_* --installroot=%s/install_root install %s"%(self.build_dir, self.build_dir, packages)) != 0:
            print "Error installing packages"
            return False

        additional_packages = []

        # now run through all .conf files and collect packages to install
        for root, dirs, files in os.walk("%s/install_root/etc/livecd"%(self.build_dir)):
            files.sort()
            for file in files:
                f = os.popen("%s/%s pkgadd"%(root, file))
                pkg_from_f = (f.read().strip()).splitlines()
                f.close()
                additional_packages += pkg_from_f

        # exclude the packages given on the command line with --exclude-package
        for e in excludePackageList:
            keep_removing = True
            while keep_removing:
                try:
                    additional_packages.remove(e)
                except ValueError:
                    keep_removing = False

        # now install the additional packages...
        packages = ""
        if len(additional_packages) > 0:
            for p in additional_packages:
                packages += " %s"%p
            if os.system("/usr/bin/yum -y -c %s/yum.conf --disablerepo=* --enablerepo=lcdr_* --installroot=%s/install_root install %s"%(self.build_dir, self.build_dir, packages)) != 0:
                print "Error installing packages"
                return False

        os.system("/bin/touch %s/install_root/etc/shadow"%(self.build_dir))
        os.system("/bin/touch %s/install_root/etc/gshadow"%(self.build_dir))

        # now run through all the post scripts from livecd scripts
        for root, dirs, files in os.walk("%s/install_root/etc/livecd"%(self.build_dir)):
            files.sort()
            for file in files:
                os.system("/usr/sbin/chroot %s/install_root /etc/livecd/%s post"%(self.build_dir, file))

        # Create initramfs
        #
        os.system("/bin/cp /usr/lib/livecd-creator/mayflower %s/install_root/sbin/mayflower"%(self.build_dir))
        os.system("/bin/cp /usr/lib/livecd-creator/run-init %s/install_root/sbin/run-init"%(self.build_dir))
        # modules needed for booting (this is butt ugly and we need to retrieve this from elsewhere, e.g. the kernel)
        mayflowerconf = open(self.build_dir + "/install_root/etc/mayflower.conf", "w")
        mayflowerconf.write('MODULES+="cdrom ide-cd ahci loop dm_snapshot squashfs ext3 ehci_hcd uhci_hcd ohci_hcd usb_storage sd_mod sr_mod usbhid ata_piix "\n')
        mayflowerconf.write('MODULES+="sata_mv sata_qstor sata_sis sata_uli"\n')
        mayflowerconf.write('MODULES+="sata_nv sata_sil24 sata_svw sata_via"\n')
        mayflowerconf.write('MODULES+="sata_promise sata_sil sata_sx4 sata_vsc"\n')
        mayflowerconf.close()

        f = os.popen("/bin/basename %s/install_root/lib/modules/*"%(self.build_dir))
        kver = f.read().strip()
        f.close()

        os.system("/usr/sbin/chroot %s/install_root /sbin/mayflower -f /boot/livecd-initramfs.img %s"%(self.build_dir, kver))
        os.system("/bin/rm -f %s/install_root/sbin/mayflower"%(self.build_dir))
        os.system("/bin/rm -f %s/install_root/sbin/run-init"%(self.build_dir))
        os.system("/bin/rm -f %s/install_root/etc/mayflower.conf"%(self.build_dir))

        # finally relabel all files
        os.system("/usr/sbin/chroot %s/install_root /sbin/fixfiles restore"%(self.build_dir))
        return True

    def configureBootloader(self):
        """configure the boot loader"""
        
        # set up boot loader
        #
        # TODO:
        #  - fix for archs not using grub
        #  - fix for non-i386
        #  - error handling
        #
        os.system("/bin/cp %s/install_root/boot/vmlinuz* %s/out/boot/vmlinuz"%(self.build_dir, self.build_dir))
        os.system("/bin/mv %s/install_root/boot/livecd-initramfs.img %s/out/boot/livecd-initramfs.img"%(self.build_dir, self.build_dir))
        os.system("/bin/cp %s/install_root/usr/share/grub/i386-redhat/stage2_eltorito %s/out/boot/grub"%(self.build_dir, self.build_dir))
        os.system("/bin/cp %s/install_root/boot/grub/splash.xpm.gz %s/out/boot/grub"%(self.build_dir, self.build_dir))

        grubconf = open(self.build_dir + "/out/boot/grub/grub.conf", "w")
        grubconf.write("default=1\n")
        grubconf.write("timeout=10\n")
        grubconf.write("splashimage=(cd)/boot/grub/splash.xpm.gz\n")
        grubconf.write("hiddenmenu\n")
        grubconf.write("title %s (Run from RAM - requires 1GB+)\n"%(self.fs_label))
        grubconf.write("        root (cd)\n")
        grubconf.write("        kernel /boot/vmlinuz ro quiet root=CDLABEL=%s rootfstype=iso9660 livecd livecd_ram\n"%(self.fs_label))
        grubconf.write("        initrd /boot/livecd-initramfs.img\n")
        grubconf.write("title %s\n"%(self.fs_label))
        grubconf.write("        root (cd)\n")
        grubconf.write("        kernel /boot/vmlinuz ro quiet root=CDLABEL=%s rootfstype=iso9660 livecd\n"%(self.fs_label))
        grubconf.write("        initrd /boot/livecd-initramfs.img\n")
        
        # TODO: enable external entitity to partipate in adding boot entries
        #
        #grubconf.write("title %s (w/ Screen Reader)\n"%(self.fs_label))
        #grubconf.write("        root (cd)\n")
        #grubconf.write("        kernel /boot/vmlinuz ro quiet root=CDLABEL=%s rootfstype=iso9660 livecd a11y_screenreader\n"%(self.fs_label))
        #grubconf.write("        initrd /boot/livecd-initramfs.img\n")
        grubconf.write("\n")
        grubconf.close()

    def createIso(self, filename):
        """write out the live CD ISO"""
        os.system("/usr/bin/mkisofs -o %s.iso -b boot/grub/stage2_eltorito -c boot/boot.catalog -no-emul-boot -boot-load-size 4 -boot-info-table -J -r -hide-rr-moved -V %s %s/out"%(filename, filename, self.build_dir))

    def createSquashFS(self):
        """create compressed squashfs file system"""
        os.system("cd %s/data; mksquashfs os.img sysroot ../out/squashfs.img"%(self.build_dir))

def usage():
    print "usage: livecd-creator [--help] "
    print "                      [--repo=<name1>,<url1> ...] [--repo=<name2>,<url2>]"
    print "                      --package=<p1> [--package=<p2> ...]"
    print "                      [--exclude-package=<e1>] --exclude-package=<e2> ...]"
    print "                      [--base-on=<path-to-iso-file>]"
    print "                      [--fslabel=<label>]"
    print "                      [--build-dir=<scratch-directory>]"
    print ""
    print " --help             : Print usage and exit"
    print " --repo             : Path to yum repository"
    print " --package          : Include this package"
    print " --exclude-package  : Exclude this package"
    print " --base-on          : Add packages to an existing live CD iso9660 image"
    print " --fslabel          : File system label (default: livecd-YYYYMMDD-HHMI)"
    print " --build-dir        : Scratch Directory (default: /var/tmp/livecd-creator/)"
    print ""

def main():

    repos = []
    packages = []
    epackages = []
    fs_label = "livecd-" + time.strftime("%Y%m%d-%H%M")
    base_on = ""
    scratch_space_dir = ""

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hr:b:p:e:f:d:",
                                   ["help", "repo=", "base-on=", "package=", "exclude-package=", "fslabel=","build-dir="])
    except getopt.GetoptError:
        print "Error..."
        print ""
        usage()
        sys.exit(2)
    output = None
    verbose = False
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit(0)
        if o in ("-r", "--repo"):
            (name, url) = a.split(",")
            for (n, u) in repos:
                if n == name:
                    print "Repo name '%s' is already in use"%n
                    sys.exit(2)
            repos.append((name, url))
            continue
        if o in ("-p", "--package"):
            packages.append(a)
            continue
        if o in ("-e", "--exclude-package"):
            epackages.append(a)
            continue
        if o in ("-f", "--fslabel"):
            fs_label = a
            continue
        if o in ("-b", "--base-on"):
            base_on = a
            continue
        if o in ("-d","--build-dir"):
            scratch_space_dir = a
            continue
        print "Unknown option %s"%o
        sys.exit(2)

    if not scratch_space_dir:
        print "No scratch directory specified.  Using default directory: /var/tmp/livecd-creator/."
        print ""
        scratch_space_dir = "/var/tmp/livecd-creator/"
        os.mkdir(scratch_space_dir)

    if len(packages) == 0:
        print "No packages specified."
        print ""
        usage()
        sys.exit(1)

    if len(repos) == 0:
        print "No repositories specified."
        print ""
        usage()
        sys.exit(1)

    target = InstallationTarget()

    if target.check_free_space(scratch_space_dir) < 3000000:
        print "Error...  Please allocate at least 3 gigabytes of free space to " + scratch_space_dir + " then retry."
        sys.exit(1)
    else:
        print "There is adequate space to build a livecd in " + scratch_space_dir + ".  Continuing with the build..."

    if not target.setup(3*1024*1024*1024, fs_label, base_on, scratch_space_dir):
        print "Cannot setup installation target. Aborting."
        sys.exit(1)

    try:
        for (n, u) in repos:
            target.addRepository(n, u)

        if not target.installPackages(packages, epackages):
            target.teardown()
            sys.exit(1)
            
        target.configureBootloader()
    except:
        target.teardown()
        print "Error during installation..."
        sys.exit(1)

    #os.system("/bin/bash")

    target.unmount()
    target.createSquashFS()
    target.createIso(fs_label)
    target.teardown()

if __name__ == "__main__":
    main()