#!/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 errno
import string
import tempfile
import time
import traceback
import subprocess
import shutil
import yum
import pykickstart.parser
import pykickstart.version

class MountError(Exception):
    def __init__(self, msg):
        Exception.__init__(self, msg)

class InstallationError(Exception):
    def __init__(self, msg):
        Exception.__init__(self, msg)

class BindChrootMount:
    """Represents a bind mount of a directory into a chroot."""
    def __init__(self, src, chroot, dest = None):
        self.src = src
        self.root = chroot

        if not dest:
            dest = src
        self.dest = self.root + "/" + dest

        self.mounted = False

    def mount(self):
        if not self.mounted:
            if not os.path.exists(self.dest):
                os.makedirs(self.dest)
            rc = subprocess.call(["/bin/mount", "--bind", self.src, self.dest])
            if rc != 0:
                raise MountError("Bind-mounting '%s' to '%s' failed" % (self.src, self.dest))
            self.mounted = True

    def umount(self):
        if self.mounted:
            rc = subprocess.call(["/bin/umount", self.dest])
            self.mounted = False
        

class LoopbackMount:
    def __init__(self, lofile, mountdir, fstype = None):
        self.lofile = lofile
        self.mountdir = mountdir
        self.fstype = fstype

        self.mounted = False
        self.losetup = False
        self.rmdir   = False
        self.loopdev = None

    def cleanup(self):
        self.umount()
        self.lounsetup()

    def umount(self):
        if self.mounted:
            rc = subprocess.call(["/bin/umount", self.mountdir])
            self.mounted = False

        if self.rmdir:
            try:
                os.rmdir(self.mountdir)
            except OSError, e:
                pass
            self.rmdir = False

    def lounsetup(self):
        if self.losetup:
            rc = subprocess.call(["/sbin/losetup", "-d", self.loopdev])
            self.losetup = False
            self.loopdev = None

    def loopsetup(self):
        if self.losetup:
            return

        rc = subprocess.call(["/sbin/losetup", "-f", self.lofile])
        if rc != 0:
            raise MountError("Failed to allocate loop device for '%s'" % self.lofile)

        # succeeded; figure out which loopdevice we're using
        buf = subprocess.Popen(["/sbin/losetup", "-a"],
                               stdout=subprocess.PIPE).communicate()[0]
        for line in buf.split("\n"):
            # loopdev: fdinfo (filename)
            fields = line.split()
            if len(fields) != 3:
                continue
            if fields[2] == "(%s)" %(self.lofile,):
                self.loopdev = fields[0][:-1]
                break

        if not self.loopdev:
            raise MountError("Failed to find loop device associated with '%s' from '/sbin/losetup -a'" % self.lofile)

        self.losetup = True

    def mount(self):
        if self.mounted:
            return

        self.loopsetup()

        if not os.path.isdir(self.mountdir):
            os.makedirs(self.mountdir)
            self.rmdir = True

        args = [ "/bin/mount", self.loopdev, self.mountdir ]
        if self.fstype:
            args.extend(["-t", self.fstype])

        rc = subprocess.call(args)
        if rc != 0:
            raise MountError("Failed to mount '%s' to '%s'" % (self.loopdev, self.mountdir))

        self.mounted = True

class SparseExt3LoopbackMount(LoopbackMount):
    def __init__(self, lofile, mountdir, size, fslabel):
        LoopbackMount.__init__(self, lofile, mountdir, fstype = "ext3")
        self.size = size
        self.fslabel = fslabel

    def _createSparseFile(self):
        dir = os.path.dirname(self.lofile)
        if not os.path.isdir(dir):
            os.makedirs(dir)

        # create the sparse file
        fd = os.open(self.lofile, os.O_WRONLY | os.O_CREAT)
        off = long(self.size * 1024L * 1024L)
        os.lseek(fd, off, 0)
        os.write(fd, '\x00')
        os.close(fd)

    def _formatFilesystem(self):
        rc = subprocess.call(["/sbin/mkfs.ext3", "-F", "-L", self.fslabel,
                              "-m", "1", self.lofile])
        if rc != 0:
            raise MountError("Error creating ext3 filesystem")
        rc = subprocess.call(["/sbin/tune2fs", "-c0", "-i0", "-Odir_index",
                              "-ouser_xattr,acl", self.lofile])

    def mount(self):
        self._createSparseFile()
        self._formatFilesystem()
        return LoopbackMount.mount(self)

class LiveCDYum(yum.YumBase):
    def __init__(self):
        yum.YumBase.__init__(self)

    def doFileLogSetup(self, uid, logfile):
        # don't do the file log for the livecd as it can lead to open fds
        # being left and an inability to clean up after ourself
        pass

    def _writeConf(self, datadir, installroot):
        conf  = "[main]\n"
        conf += "installroot=%s\n" % installroot
        conf += "cachedir=/var/cache/yum\n"
        conf += "plugins=0\n"
        conf += "reposdir=\n"

        path = datadir + "/yum.conf"

        f = file(path, "w+")
        f.write(conf)
        f.close()

        os.chmod(path, 0644)

        return path

    def setup(self, datadir, installroot):
        self.doConfigSetup(fn = self._writeConf(datadir, installroot),
                           root = installroot)
        self.conf.cache = 0
        self.doTsSetup()
        self.doRpmDBSetup()
        self.doRepoSetup()
        self.doSackSetup()

    def selectPackage(self, pkg):
        """Select a given package.  Can be specified with name.arch or name*"""
        try:
            self.install(pattern = pkg)
        except yum.Errors.InstallError, e:
            raise InstallationError("Failed to find package '%s' : %s" % (pkg, e))
        
    def deselectPackage(self, pkg):
        """Deselect package.  Can be specified as name.arch or name*"""
        sp = pkg.rsplit(".", 2)
        txmbrs = []
        if len(sp) == 2:
            txmbrs = self.tsInfo.matchNaevr(name=sp[0], arch=sp[1])

        if len(txmbrs) == 0:
            exact, match, unmatch = yum.packages.parsePackages(self.pkgSack.returnPackages(), [pkg], casematch=1)
            for p in exact + match:
                txmbrs.append(p)

        if len(txmbrs) > 0:
            map(lambda x: self.tsInfo.remove(x.pkgtup), txmbrs)
        else:
            print >> sys.stderr, "No such package %s to remove" %(pkg,)

    def selectGroup(self, grp, include = pykickstart.parser.GROUP_DEFAULT):
        try:
            yum.YumBase.selectGroup(self, grp)
        except (yum.Errors.InstallError, yum.Errors.GroupsError), e:
            raise InstallationError("Failed to find group '%s' : %s" % (grp, e))
        if include == pykickstart.parser.GROUP_REQUIRED:
            map(lambda p: self.deselectPackage(p), grp.default_packages.keys())
        elif include == pykickstart.parser.GROUP_ALL:
            map(lambda p: self.selectPackage(p), grp.optional_packages.keys())
        
    def addRepository(self, name, url = None, mirrorlist = None):
        repo = yum.yumRepo.YumRepository(name)
        if url:
            repo.baseurl.append(url)
        if mirrorlist:
            repo.mirrorlist = mirrorlist
        conf = yum.config.RepoConf()
        for k, v in conf.iteritems():
            if v or not hasattr(repo, k):
                repo.setAttribute(k, v)
        repo.basecachedir = self.conf.cachedir
        repo.metadata_expire = 0
        # disable gpg check???
        repo.gpgcheck = 0
        repo.enable()
        repo.setup(0)
        self.repos.add(repo)
            
    def runInstall(self):
        (res, resmsg) = self.buildTransaction()
        if res != 2:
            raise InstallationError("Failed to build transaction : %s" % str.join("\n", resmsg))
        
        dlpkgs = map(lambda x: x.po, filter(lambda txmbr: txmbr.ts_state in ("i", "u"), self.tsInfo.getMembers()))
        self.downloadPkgs(dlpkgs)
        # FIXME: sigcheck?
        
        self.initActionTs()
        self.populateTs(keepold=0)
        self.ts.check()
        self.ts.order()
        # FIXME: callback should be refactored a little in yum 
        sys.path.append('/usr/share/yum-cli')
        import callback
        cb = callback.RPMInstallCallback()
        cb.tsInfo = self.tsInfo
        cb.filelog = False
        return self.runTransaction(cb)

class InstallationTarget:
    def __init__(self, repos, packages, epackages, groups, fs_label, skip_compression, skip_prelink):
        self.ayum = None
        self.repos = repos
        self.packages = packages
        self.epackages = epackages
        self.groups = groups
        self.fs_label = fs_label
        self.skip_compression = skip_compression
        self.skip_prelink = skip_prelink

        self.build_dir = None
        self.instloop = None
        self.bindmounts = []
        self.ksparser = None
        
    def parse(self, kscfg):
        ksversion = pykickstart.version.makeVersion()
        self.ksparser = pykickstart.parser.KickstartParser(ksversion)
        if kscfg:
            try:
                self.ksparser.readKickstart(kscfg)
            except IOError, (err, msg):
                raise InstallationError("Failed to read kickstart file '%s' : %s" % (kscfg, msg))
            except pykickstart.errors.KickstartError, e:
                raise InstallationError("Failed to parse kickstart file '%s' : %s" % (kscfg, e))

            for repo in self.ksparser.handler.repo.repoList:
                already_given = False
                for cmd_name, cmd_url in self.repos:
                    if cmd_name == repo.name:
                        already_given = True
                        break

                if not already_given:
                    self.repos.append( (repo.name, repo.baseurl) )

            self.packages.extend(self.ksparser.handler.packages.packageList)
            self.groups.extend(map(lambda g: (g.name, g.include),
                                   self.ksparser.handler.packages.groupList))
            self.epackages.extend(self.ksparser.handler.packages.excludedList)

        if not self.packages and not self.groups:
            raise InstallationError("No packages or groups specified")

        if not self.repos:
            raise InstallationError("No repositories specified")

    def base_on_iso(self, base_on):
        """helper function to extract ext3 file system from a live CD ISO"""

        isoloop = LoopbackMount(base_on, "%s/base_on_iso" %(self.build_dir,))

        squashloop = LoopbackMount("%s/squashfs.img" %(isoloop.mountdir,),
                                   "%s/base_on_squashfs" %(self.build_dir,),
                                   "squashfs")

        try:
            try:
                isoloop.mount()
            except MountError, e:
                raise InstallationError("Failed to loopback mount '%s' : %s" % (base_on, e))

            if not os.path.exists(squashloop.lofile):
                raise InstallationError("'%s' is not a valid live CD ISO : squashfs.img doesn't exist" % base_on)

            try:
                squashloop.mount()
            except MountError, e:
                raise InstallationError("Failed to loopback mount squashfs.img from '%s' : %s" % (base_on, e))

            os_image = self.build_dir + "/base_on_squashfs/os.img"

            if not os.path.exists(os_image):
                raise InstallationError("'%s' is not a valid live CD ISO : os.img doesn't exist" % base_on)

            shutil.copyfile(os_image, self.build_dir + "/data/os.img")
        finally:
            # unmount and tear down the mount points and loop devices used
            squashloop.cleanup()
            isoloop.cleanup()

    def write_fstab(self):
        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()

    def setup(self, image_size, base_on = None):
        """setup target ext3 file system in preparation for an install"""

        # setup temporary build dirs
        try:
            self.build_dir = tempfile.mkdtemp(dir="/var/tmp", prefix="livecd-creator-")
        except OSError, (err, msg):
            raise InstallationError("Failed create build directory in /var/tmp: %s" % msg)

        os.makedirs(self.build_dir + "/out/isolinux")
        os.makedirs(self.build_dir + "/out/sysroot")
        os.makedirs(self.build_dir + "/data/sysroot")
        os.makedirs(self.build_dir + "/install_root")
        os.makedirs(self.build_dir + "/yum-cache")

        if base_on:
            # get backing ext3 image if we're based this build on an existing live CD ISO
            self.base_on_iso(base_on)

            self.instloop = LoopbackMount("%s/data/os.img" %(self.build_dir,),
                                          "%s/install_root" %(self.build_dir,))
        else:
            self.instloop = SparseExt3LoopbackMount("%s/data/os.img"
                                                    %(self.build_dir,),
                                                    "%s/install_root"
                                                    %(self.build_dir,),
                                                    image_size,
                                                    self.fs_label)
            

        try:
            self.instloop.mount()
        except MountError, e:
            raise InstallationError("Failed to loopback mount '%s' : %s" % (self.instloop.lofile, e))

        if not base_on:
            # create a few directories needed if it's a new image
            os.makedirs(self.build_dir + "/install_root/etc")
            os.makedirs(self.build_dir + "/install_root/boot")
            os.makedirs(self.build_dir + "/install_root/var/log")
            os.makedirs(self.build_dir + "/install_root/var/cache/yum")

        # bind mount system directories into install_root/
        for (f, dest) in [("/sys", None), ("/proc", None), ("/dev", None),
                          ("/dev/pts", None), ("/selinux", None),
                          (self.build_dir + "/yum-cache", "/var/cache/yum")]:
            self.bindmounts.append(BindChrootMount(f, self.build_dir + "/install_root", dest))

        for b in self.bindmounts:
            b.mount()

        # make sure /etc/mtab is current inside install_root
        os.symlink("../proc/mounts", self.build_dir + "/install_root/etc/mtab")

        self.write_fstab()

        self.ayum = LiveCDYum()
        self.ayum.setup(self.build_dir + "/data",
                        self.build_dir + "/install_root")

    def unmount(self):
        """detaches system bind mounts and install_root for the file system and tears down loop devices used"""
        self.ayum = None

        # FIXME: Make one last ditch effort to close fds still open
        # in the install root; this is only needed because
        # there's no way to ask yum to close its sqlite dbs,
        # though. See https://bugzilla.redhat.com/236409
        for i in range(3, os.sysconf("SC_OPEN_MAX")):
            try:
                os.close(i)
            except:
                pass

        try:
            os.unlink(self.build_dir + "/install_root/etc/mtab")
        except OSError:
            pass

        self.bindmounts.reverse()
        for b in self.bindmounts:
            b.umount()

        if self.instloop:
            self.instloop.cleanup()
            self.instloop = None

    def teardown(self):
        if self.build_dir:
            self.unmount()
            shutil.rmtree(self.build_dir, ignore_errors = True)

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

    def run_in_root(self):
        os.chroot("%s/install_root" %(self.build_dir,))
        os.chdir("/")

    def installPackages(self, packageList, excludePackageList, groupList = []):
        """install packages into target file system"""
        try:
            map(lambda pkg: self.ayum.selectPackage(pkg), packageList)
            map(lambda grp: self.ayum.selectGroup(grp[0], grp[1]), groupList)
            map(lambda pkg: self.ayum.deselectPackage(pkg), excludePackageList)

            self.ayum.runInstall()
        finally:
            self.ayum.closeRpmDB()

    def writeNetworkIfCfg(self, instroot, network):
        path = instroot + "/etc/sysconfig/network-scripts/ifcfg-" + network.device

        f = file(path, "w+")
        os.chmod(path, 0644)

        f.write("DEVICE=%s\n" % network.device)
        f.write("BOOTPROTO=%s\n" % network.bootProto)

        if network.bootProto.lower() == "static":
            if network.ip:
                f.write("IPADDR=%s\n" % network.ip)
            if network.netmask:
                f.write("NETMASK=%s\n" % network.netmask)

        if network.onboot:
            f.write("ONBOOT=on\n")
        else:
            f.write("ONBOOT=off\n")

        if network.essid:
            f.write("ESSID=%s\n" % network.essid)

        if network.ethtool:
            if network.ethtool.find("autoneg") == -1:
                network.ethtool = "autoneg off " + network.ethtool
            f.write("ETHTOOL_OPTS=%s\n" % network.ethtool)

        if network.bootProto.lower() == "dhcp":
            if network.hostname:
                f.write("DHCP_HOSTNAME=%s\n" % network.hostname)
            if network.dhcpclass:
                f.write("DHCP_CLASSID=%s\n" % network.dhcpclass)

        if network.mtu:
            f.write("MTU=%s\n" % network.mtu)

        f.close()

    def writeNetworkKey(self, instroot, network):
        if not network.wepkey:
            return

        path = instroot + "/etc/sysconfig/network-scripts/keys-" + network.device
        f = file(path, "w+")
        os.chmod(path, 0600)
        f.write("KEY=%s\n" % network.wepkey)
        f.close()

    def writeNetworkConfig(self, instroot, useipv6, hostname, gateway):
        path = instroot + "/etc/sysconfig/network"
        f = file(path, "w+")
        os.chmod(path, 0644)

        f.write("NETWORKING=yes\n")

        if useipv6:
            f.write("NETWORKING_IPV6=yes\n")
        else:
            f.write("NETWORKING_IPV6=no\n")

        if hostname:
            f.write("HOSTNAME=%s\n" % hostname)
        else:
            f.write("HOSTNAME=localhost.localdomain\n")

        if gateway:
            f.write("GATEWAY=%s\n" % gateway)

        f.close()

    def writeNetworkHosts(self, instroot, hostname):
        localline = ""
        if hostname and hostname != "localhost.localdomain":
            localline += hostname + " "
            l = string.split(hostname, ".")
            if len(l) > 1:
                localline += l[0] + " "
        localline += "localhost.localdomain localhost"

        path = instroot + "/etc/hosts"
        f = file(path, "w+")
        os.chmod(path, 0644)
        f.write("127.0.0.1\t\t%s\n" % localline)
        f.write("::1\t\tlocalhost6.localdomain6 localhost6\n")
        f.close()

    def writeNetworkResolv(self, instroot, nodns, primaryns, secondaryns):
        if nodns or not primaryns:
            return

        path = instroot + "/etc/resolv.conf"
        f = file(path, "w+")
        os.chmod(path, 0644)

        for ns in (primaryns, secondaryns):
            if ns:
                f.write("nameserver %s\n" % ns)

        f.close()

    def configureNetwork(self):
        instroot = self.build_dir + "/install_root"

        try:
            os.makedirs(instroot + "/etc/sysconfig/network-scripts")
        except OSError, (err, msg):
            if err != errno.EEXIST:
                raise

        useipv6 = False
        nodns = False
        hostname = None
        gateway = None
        primaryns = None
        secondaryns = None

        for network in self.ksparser.handler.network.network:
            if not network.device:
                raise InstallationError("No --device specified with network kickstart command")

            if network.onboot and network.bootProto.lower() != "dhcp" and \
               not (network.ip and network.netmask):
                raise InstallationError("No IP address and/or netmask specified with static " +
                                        "configuration for '%s'" % network.device)

            self.writeNetworkIfCfg(instroot, network)
            self.writeNetworkKey(instroot, network)

            if network.ipv6:
                useipv6 = True
            if network.nodns:
                nodns = True

            if network.hostname:
                hostname = network.hostname
            if network.gateway:
                gateway = network.gateway

            if network.nameserver:
                nameservers = string.split(network.nameserver, ",")
                if len(nameservers) >= 1:
                    primaryns = nameservers[0]
                if len(nameservers) >= 2:
                    secondayns = nameservers[0]

        self.writeNetworkConfig(instroot, useipv6, hostname, gateway)
        self.writeNetworkHosts(instroot, hostname)
        self.writeNetworkResolv(instroot, nodns, primaryns, secondaryns)

    def configureSystem(self):
        instroot = "%s/install_root" %(self.build_dir,)
        
        # FIXME: this is a bit ugly, but with the current pykickstart
        # API, we don't really have a lot of choice.  it'd be nice to
        # be able to do something different, but so it goes

        # set up the language
        lang = self.ksparser.handler.lang.lang or "en_US.UTF-8"
        f = open("%s/etc/sysconfig/i18n" %(instroot,), "w+")
        f.write("LANG=\"%s\"\n" %(lang,))
        f.close()

        # next, the keyboard
        # FIXME: should this impact the X keyboard config too???
        # or do we want to make X be able to do this mapping
        import rhpl.keyboard
        k = rhpl.keyboard.Keyboard()
        if self.ksparser.handler.keyboard.keyboard:
            k.set(self.ksparser.handler.keyboard.keyboard)
        k.write(instroot)

        # next up is timezone
        tz = self.ksparser.handler.timezone.timezone or "America/New_York"
        utc = self.ksparser.handler.timezone.isUtc
        f = open("%s/etc/sysconfig/clock" %(instroot,), "w+")
        f.write("ZONE=\"%s\"\n" %(tz,))
        f.write("UTC=%s\n" %(utc,))
        f.close()

        # do any authconfig bits
        auth = self.ksparser.handler.authconfig.authconfig or "--useshadow --enablemd5"
        if os.path.exists("%s/usr/sbin/authconfig" %(instroot,)):
            args = ["/usr/sbin/authconfig", "--update", "--nostart"]
            args.extend(auth.split())
            subprocess.call(args, preexec_fn=self.run_in_root)

        # firewall.  FIXME: should handle the rest of the options
        if self.ksparser.handler.firewall.enabled and os.path.exists("%s/usr/sbin/lokkit" %(instroot,)):
            subprocess.call(["/usr/sbin/lokkit", "-f", "--quiet",
                             "--nostart", "--enabled"],
                            preexec_fn=self.run_in_root)

        # selinux
        if os.path.exists("%s/usr/sbin/lokkit" %(instroot,)):
            args = ["/usr/sbin/lokkit", "-f", "--quiet", "--nostart"]
            if self.ksparser.handler.selinux.selinux:
                args.append("--selinux=enforcing")
            else:
                args.append("--selinux=disabled")
            subprocess.call(args, preexec_fn=self.run_in_root)

        # Set the root password
        if self.ksparser.handler.rootpw.isCrypted:
            subprocess.call(["/usr/sbin/usermod", "-p", self.ksparser.handler.rootpw.password, "root"], preexec_fn=self.run_in_root)
        elif self.ksparser.handler.rootpw.password == "":
            # Root password is not set and not crypted, empty it
            subprocess.call(["/usr/bin/passwd", "-d", "root"], preexec_fn=self.run_in_root)
        else:
            # Root password is set and not crypted
            p1 = subprocess.Popen(["/bin/echo", self.ksparser.handler.rootpw.password], stdout=subprocess.PIPE, preexec_fn=self.run_in_root)
            p2 = subprocess.Popen(["/usr/bin/passwd", "--stdin", "root"], stdin=p1.stdout, stdout=subprocess.PIPE, preexec_fn=self.run_in_root)
            output = p2.communicate()[0]

        # enable/disable services appropriately
        if os.path.exists("%s/sbin/chkconfig" %(instroot,)):
            for s in self.ksparser.handler.services.enabled:
                subprocess.call(["/sbin/chkconfig", s, "--level", "345", "on"],
                                preexec_fn=self.run_in_root)
            for s in self.ksparser.handler.services.disabled:
                subprocess.call(["/sbin/chkconfig", s, "--level", "345", "off"],
                                preexec_fn=self.run_in_root)

        # x by default?
        if self.ksparser.handler.xconfig.startX:
            f = open("%s/etc/inittab" %(instroot,), "rw+")
            buf = f.read()
            buf = buf.replace("id:3:initdefault", "id:5:initdefault")
            f.seek(0)
            f.write(buf)
            f.close()

        # and now, for arbitrary %post scripts
        for s in filter(lambda s: s.type == pykickstart.parser.KS_SCRIPT_POST,
                        self.ksparser.handler.scripts):
            # we can only safely run scripts in the chroot
            if not s.inChroot:
                print >> sys.stderr, "Not running script outside of chroot"
                continue

            (fd, path) = tempfile.mkstemp("", "ks-script-", "%s/tmp" %(instroot,))
            os.write(fd, s.script)
            os.close(fd)
            os.chmod(path, 0700)

            try:
                subprocess.call([s.interp, "/tmp/%s" %(os.path.basename(path),)],
                                preexec_fn = self.run_in_root)
            except OSError, (err, msg):
                os.unlink(path)
                raise InstallationError("Failed to execute %%post script with '%s' : %s" % (s.interp, msg))
            os.unlink(path)

    def get_kernel_version(self):
        #
        # FIXME: this doesn't handle multiple kernels - we should list
        #        them all in the isolinux menu
        #
        kernels = []
        modules_dir = "%s/install_root/lib/modules" % self.build_dir

        if os.path.isdir(modules_dir):
            kernels = os.listdir(modules_dir)

        if not kernels:
            raise InstallationError("No kernels installed: /lib/modules is empty")

        return kernels[0]

    def createInitramfs(self):
        # Create initramfs
        if not os.path.isfile("/usr/lib/livecd-creator/mayflower"):
            raise InstallationError("livecd-creator not correctly installed : "+
                                    "/usr/lib/livecd-creator/mayflower not found")
        shutil.copy("/usr/lib/livecd-creator/mayflower",
                        "%s/install_root/sbin/mayflower" %(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 ext2 ehci_hcd uhci_hcd ohci_hcd usb_storage sd_mod sr_mod usbhid ata_piix vfat msdos "\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.write('MODULES+="ata_generic pata_ali pata_amd pata_artop pata_atiixp pata_cmd64x pata_cs5520 pata_cs5530 pata_cs5535 pata_cypress pata_efar pata_hpt366 pata_hpt37x pata_hpt3x2n pata_hpt3x3 pata_isapnp pata_it821x pata_jmicron pata_marvell pata_mpiix pata_netcell pata_ns87410 pata_oldpiix pata_optidma pata_opti pata_pcmcia pata_pdc2027x pata_pdc202xx_old pata_qdi pata_serverworks pata_sil680 pata_sis pata_sl82c105 pata_triflex pata_via pdc_adma "\n')
        mayflowerconf.write('MODULES+="sym53c8xx aic7xxx "\n')        
        mayflowerconf.close()

        subprocess.call(["/sbin/mayflower", "-f", "/boot/livecd-initramfs.img",
                        self.get_kernel_version()],
                        preexec_fn=self.run_in_root)
        for f in ("/sbin/mayflower", "/etc/mayflower.conf"):
            os.unlink("%s/install_root/%s" %(self.build_dir, f))

    def relabelSystem(self):
        # finally relabel all files
        if self.ksparser.handler.selinux.selinux:
            instroot = "%s/install_root" %(self.build_dir,)
            if os.path.exists("%s/sbin/restorecon" %(instroot,)):
                subprocess.call(["/sbin/restorecon", "-v", "-r", "/"],
                                preexec_fn=self.run_in_root)

    def prelinkSystem(self):
        # prelink the system
        instroot = "%s/install_root" %(self.build_dir,)
        if os.path.exists("%s/usr/sbin/prelink" %(instroot,)):
            subprocess.call(["/usr/sbin/prelink", "-mRaN"],
                            preexec_fn=self.run_in_root)
        return True

    def launchShell(self):
        subprocess.call(["/bin/bash"], preexec_fn=self.run_in_root)

    def configureBootloader(self):
        """configure the boot loader"""
        
        # set up boot loader
        #
        # TODO:
        #  - fix for archs not using grub
        #  - fix for non-i386
        #  - error handling
        #
        shutil.copyfile("%s/install_root/boot/vmlinuz-%s"
                        %(self.build_dir, self.get_kernel_version()),
                        "%s/out/isolinux/vmlinuz" %(self.build_dir,))

        shutil.copyfile("%s/install_root/boot/livecd-initramfs.img"
                        %(self.build_dir,),
                        "%s/out/isolinux/initrd.img" %(self.build_dir,))
        os.unlink("%s/install_root/boot/livecd-initramfs.img"
                  %(self.build_dir,))

        for p in ["isolinux.bin", "vesamenu.c32"]:
            path = "%s/install_root/usr/lib/syslinux/%s" % (self.build_dir, p)

            if not os.path.isfile(path):
                raise InstallationError("syslinux not installed : %s not found" % path)

            shutil.copy(path, "%s/out/isolinux/%s" % (self.build_dir, p))

        if os.path.exists("%s/install_root/usr/lib/anaconda-runtime/syslinux-vesa-splash.jpg" %(self.build_dir,)):
            shutil.copy("%s/install_root/usr/lib/anaconda-runtime/syslinux-vesa-splash.jpg" %(self.build_dir,),
                        "%s/out/isolinux/splash.jpg" %(self.build_dir,))
            have_background = "menu background splash.jpg"
        else:
            have_background = ""

        cfg = """
default vesamenu.c32
timeout 600

%(background)s
menu title Welcome to %(label)s!
menu color border 0 #ffffffff #00000000
menu color sel 0 #ffffffff #ff000000
menu color title 0 #ffffffff #00000000
menu color tabmsg 0 #ffffffff #00000000
menu color unsel 0 #ffffffff #00000000
menu color hotsel 0 #ff000000 #ffffffff
menu color hotkey 0 #ffffffff #ff000000
""" %{"label": self.fs_label, "background" : have_background}

        stanzas = [("linux", "Run from image", ""),
                   ("runfromram", "Run from RAM - requires 1 GB+", "live_ram")]
        if os.path.exists("%s/install_root/usr/lib/anaconda-runtime/checkisomd5" %(self.build_dir,)):
            stanzas.append( ("check", "Verify and run from image", "check") )

        for (short, long, extra) in stanzas:
            cfg += """label %(short)s
  menu label %(long)s
  kernel vmlinuz
  append initrd=initrd.img  ro quiet root=CDLABEL=%(label)s rootfstype=iso9660 liveimg %(extra)s
""" %{"label": self.fs_label, "background" : have_background,
      "short": short, "long": long, "extra": extra}

        cfgf = open("%s/out/isolinux/isolinux.cfg" %(self.build_dir,), "w")
        cfgf.write(cfg)
        cfgf.close()
        
        # TODO: enable external entitity to partipate in adding boot entries

    def install(self):
        for (name, url) in self.repos:
            self.ayum.addRepository(name, url)

        self.installPackages(self.packages, self.epackages, self.groups)
        self.configureSystem()
        self.configureNetwork()
        self.relabelSystem()
        if not self.skip_prelink:
            self.prelinkSystem()
        self.createInitramfs()
        self.configureBootloader()

    def createIso(self):
        """write out the live CD ISO"""
        subprocess.call(["/usr/bin/mkisofs", "-o", "%s.iso" %(self.fs_label,),
                         "-b", "isolinux/isolinux.bin",
                         "-c", "isolinux/boot.cat",
                         "-no-emul-boot", "-boot-load-size", "4",
                         "-boot-info-table", "-J", "-r", "-hide-rr-moved",
                         "-V", "%s" %(self.fs_label,), "%s/out" %(self.build_dir)])

        # implant an isomd5sum
        if os.path.exists("/usr/lib/anaconda-runtime/implantisomd5"):
            subprocess.call(["/usr/lib/anaconda-runtime/implantisomd5",
                             "%s.iso" %(self.fs_label,)])
        else:
            print >> sys.stderr, "anaconda-runtime not installed; not setting up mediacheck"

    def createSquashFS(self):
        """create compressed squashfs file system"""
        if not self.skip_compression:
            # FIXME: mksquashfs segfaults if PWD isn't set in the environment
            subprocess.call(["/sbin/mksquashfs", "os.img", "sysroot",
                             "../out/squashfs.img"],
                            cwd="%s/data" %(self.build_dir,),
                            env={"PWD": "%s/data" %(self.build_dir,)})
        else:
            shutil.move("%s/data/os.img" %(self.build_dir,),
                        "%s/out/ext3fs.img" %(self.build_dir,))

    def package(self):
        self.createSquashFS()
        self.createIso()

def usage(out):
    print >> out, """
usage: livecd-creator [--help] 
                      [--config=<path-to-kickstart-file> | --repo=<name>,<url> --package=<p>]
                      [--repo=<name1>,<url1>] [--repo=<name2>,<url2> ...]
                      [--package=<p1>] [--package=<p2> ...]
                      [--exclude-package=<e1>] [--exclude-package=<e2> ...]
                      [--base-on=<path-to-iso-file>]
                      [--fslabel=<label>]
                      [--skip-compression]
                      [--uncompressed-size=<size-in-MB>]
                      [--shell]

 --help              : Print usage and exit
 --config            : Path to kickstart config file
 --repo              : Add / override yum repository
 --package           : Include this package
 --exclude-package   : Exclude this package
 --base-on           : Add packages to an existing live CD iso9660 image
 --fslabel           : File system label (default: livecd-YYYYMMDD-HHMI)
 --skip-compression  : Don't compress the image
 --prelink           : Prelink the image
 --uncompressed-size : Size of uncompressed fs in MB (default: 4096)
 --shell             : Start a shell in the chroot for post-configuration

 Examples:

  Create minimal live cd:
  # livecd-creator --config=/usr/share/livecd-tools/livecd-fedora-minimal.ks

  Create minimal live cd and use local package repositories:
  # livecd-creator --config=/usr/share/livecd-tools/livecd-fedora-minimal.ks \\
                   --repo=a-dev,file:///home/user/core/RPMS             \\
                   --repo=a-extras-dev,file:///home/user/extras/RPMS
"""

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

class Options:
    def __init__(self):
        self.repos = []
        self.packages = []
        self.groups = []
        self.epackages = []
        self.fs_label = "livecd-" + time.strftime("%Y%m%d-%H%M")
        self.base_on = None
        self.kscfg = None
        self.skip_compression = False
        self.skip_prelink = True
        self.uncompressed_size = 4096
        self.give_shell = False

def parse_options(args):
    try:
        opts, args = getopt.getopt(args, "hr:b:p:e:f:c:su:l",
                                   ["help", "repo=", "base-on=", "package=",
                                    "exclude-package=", "fslabel=", "config=",
                                    "skip-compression", "uncompressed-size=",
                                    "shell", "no-prelink", "prelink"])
    except getopt.GetoptError, msg:
        raise Usage(msg)

    options = Options()

    for o, a in opts:
        if o in ("-h", "--help"):
            raise Usage(no_error = True)
        if o in ("-l", "--shell"):
            options.give_shell = True
            continue
        if o in ("-s", "--skip-compression"):
            options.skip_compression = True
            continue
        if o in ("--no-prelink",):
            options.skip_prelink = True
            continue
        if o in ("--prelink",):
            options.skip_prelink = False
            continue
        if o in ("-u", "--uncompressed-size"):
            options.uncompressed_size = int(a)
            continue
        if o in ("-c", "--config"):
            options.kscfg = a
            if not os.path.isfile(options.kscfg):
                raise Usage("Kickstart config '%s' does not exist" % options.kscfg)
            continue
        if o in ("-r", "--repo"):
            (name, url) = a.split(",")
            for (n, u) in options.repos:
                if n == name:
                    raise Usage("Repo name '%s' is already in use" % n)
            options.repos.append((name, url))
            continue
        if o in ("-p", "--package"):
            if a.startswith("@"):
                options.groups.append((a[1:], pykickstart.parser.GROUP_DEFAULT))
            else:
                options.packages.append(a)
            continue
        if o in ("-e", "--exclude-package"):
            options.epackages.append(a)
            continue
        if o in ("-f", "--fslabel"):
            options.fs_label = a
            if len(options.fs_label) > 32:
                raise Usage("CD labels are limited to 32 characters")
            continue
        if o in ("-b", "--base-on"):
            options.base_on = a
            if not os.path.isfile(options.base_on):
                raise Usage("Live CD ISO '%s' does not exist" % options.base_on)
            continue
        raise Usage("Unknown option %s" % o)

    if not options.kscfg and not (options.packages or options.groups):
        raise Usage("No packages or groups specified")

    if not options.kscfg and not options.repos:
        raise Usage("No repositories specified")

    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
        usage(out)
        return ret

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

    target = InstallationTarget(options.repos,
                                options.packages,
                                options.epackages,
                                options.groups,
                                options.fs_label,
                                options.skip_compression,
                                options.skip_prelink)

    try:
        target.parse(options.kscfg)

        target.setup(options.uncompressed_size, options.base_on)

        target.install()

        if options.give_shell:
            print "Launching shell. Exit to continue."
            print "----------------------------------"
            target.launchShell()

        target.unmount()

        target.package()
    except InstallationError, e:
        print >> sys.stderr, "Error creating Live CD : %s" % e
        target.teardown()
        return 1
    target.teardown()

    return 0

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