#!/usr/bin/env python
# Copyright (C) 2007 Platform Computing Inc
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of version 2 of the GNU General Public License as
# published by the Free Software Foundation.
# 	
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
#
# $Id: mk-bootable 1016 2008-03-27 12:26:15Z ggoh $
#

import os
import sys
import subprocess
import tempfile
import atexit
import shutil
import optparse
import glob

class BuildProfile:
    
    def __init__(self, **kwargs):
        self.instsource = kwargs.get('instsource',None)
        self.distro = kwargs.get('distro',None)
        self.version = kwargs.get('version',None)
        self.arch = kwargs.get('arch',None)
        self.kitsource = kwargs.get('kitsource',None)


def cleanup(tdir):
    """Housekeeping routines"""
    if os.path.exists(tdir):
        shutil.rmtree(tdir)
        
def getKusuTrunk():
    """Returns the SVN trunk for Kusu/Emerald"""
    script = sys.argv[0]
    bindir = os.path.dirname(script)
    trunk = os.path.normpath(bindir + '/..')
    return trunk
    
def patchInitrd(initrd, kscfg):
    """Patches the initrd"""
    
    # create temp scratch dir
    scratchdir = tempfile.mkdtemp(prefix='initrd-')
    initrddir = tempfile.mkdtemp(dir=scratchdir)
    pinitrd = os.path.join(scratchdir,'initrd.img')
    
    # unpack the initrd
    # FIXME: only supporting RHEL 5 and CentOS 5 for now..
    exitcode = unpackRHELInitrd(initrd, initrddir)
    
    # copy in the ks.cfg
    shutil.copy(kscfg,os.path.join(initrddir,'ks.cfg'))
    
    # repacking the initrd
    packInitramFS(initrddir,pinitrd)
    
    # remove the original initrd and replace it with the patched version
    os.remove(initrd)
    shutil.move(pinitrd,initrd)
    
    shutil.rmtree(scratchdir)

    
def packInitramFS(dirname, rootimgpath):
    """Packs the directory dir into an initramfs rootimgpath. The initscript template file can be passed as well.
       Refer to Linux Kernel Documentation/filesystems/ramfs-rootfs-initramfs.txt for more information.
    """

    if not os.path.exists(dirname): 
        print 'Directory %s does not exist!' % dirname
        sys.exit(-1)

    if os.getuid() <> 0: 
        print 'Operation requires root access!'
        sys.exit(-1)

    # nuke existing rootimgpath if existing
    if os.path.exists(rootimgpath): os.remove(rootimgpath)

    cwd = os.path.abspath(dirname)
    findP = subprocess.Popen("find .",cwd=cwd,shell=True,stdout=subprocess.PIPE)
    cpioP = subprocess.Popen("cpio -o -H newc --quiet",cwd=cwd,shell=True,stdin=findP.stdout,stdout=subprocess.PIPE)
    gzipP = subprocess.Popen("gzip > %s" % rootimgpath,shell=True,stdin=cpioP.stdout)
    gzipP.communicate()

    return gzipP.returncode

def unpackRHELInitrd(rootimgpath, dirname):
    """Unpacks the file rootimgpath into the directory dir."""

    if not os.path.exists(rootimgpath):
        print '%s does not exist!' % rootimgpath
        sys.exit(-1)

    if os.getuid() <> 0:
        print 'Operation requires root access!'
        sys.exit(-1)

    # nuke existing dirname
    dirname = os.path.abspath(dirname)
    if os.path.exists(dirname): shutil.rmtree(dirname)
    os.makedirs(dirname)

    zcatP = subprocess.Popen('zcat %s' % rootimgpath,shell=True,stdout=subprocess.PIPE)
    cpioP = subprocess.Popen('cpio --quiet -id',cwd=dirname,shell=True,stdin=zcatP.stdout,stdout=subprocess.PIPE)
    cpioP.communicate()
    
    return cpioP.returncode


unpackCentOSInitrd = unpackRHELInitrd

def makeSuseBootDir(rootdir, buildprofile):
    """Creates a boot directory in rootdir with prefilled settings for the distro"""
    # look for the distro's boot dir
    bootdir = os.path.join(buildprofile.instsource,'boot')
    if not os.path.exists(bootdir):
        print 'No boot directory found in %s !' % buildprofile.instsource
        sys.exit(-1)
    print 'Found %s ..' % bootdir

    # copy the contents over to the newly made ext2fs filesystem
    #cmd = '(find boot | cpio -mpdu %s) >/dev/null 2>&1'
    cmd = "(find * | grep -v 'suse/' | grep -v 'docu/' | grep -v 'dosutils/' | grep -v 'patches/' | grep -v '\.jpg' | grep -v '\.tr' | grep -v '\.pcx' | grep -v '\.hlp' | grep -v '\.spl' | cpio -mpdu %s) >/dev/null 2>&1"
    cpioP = subprocess.Popen(cmd % rootdir,shell=True,cwd=buildprofile.instsource,stdout=subprocess.PIPE)
    cpioP.wait()
    
    bootloaddir = os.path.join(rootdir,'boot/%s/loader' % buildprofile.arch)
    if not os.path.exists(bootdir):
        print 'Copy from %s to %s failed!' % ( bootdir, rootdir)
        sys.exit(-1)
    print 'Using %s ..' % bootdir
    
    rootimg = os.path.join(rootdir,'boot/%s/root' % buildprofile.arch)
    if not os.path.exists(rootimg):
        print 'No root img found in %s/boot/%s !' % (rootdir, buildprofile.arch)
        sys.exit(-1)
    print 'Using %s as the root image..' % rootimg
    
    # get the svn trunk layout
    trunk = getKusuTrunk()
    
    # grab the isolinux.cfg
    isolinuxcfg = os.path.join(trunk, \
                    'src/dists/%(distro)s/%(version)s/%(arch)s/initrd/isolinux.cfg' %
                    {'distro':buildprofile.distro,
                    'version':buildprofile.version,
                    'arch':buildprofile.arch})                    
    if not os.path.exists(isolinuxcfg):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % isolinuxcfg
        sys.exit(-1)

    print 'Using %s as the isolinux.cfg file..' % isolinuxcfg
    
    shutil.copy(isolinuxcfg,bootloaddir)
        
    # grab the boot.msg
    bootmsg = os.path.join(trunk, \
                    'src/dists/%(distro)s/%(version)s/%(arch)s/boot.msg' %
                    {'distro':buildprofile.distro,
                    'version':buildprofile.version,
                    'arch':buildprofile.arch})
    print 'Using %s.' % bootmsg
    if not os.path.exists(bootmsg):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % bootmsg
        sys.exit(-1)
    
    shutil.copy(bootmsg,bootloaddir)
    
    # grab the splash.lss
    splashlss = os.path.join(trunk, \
                    'src/dists/common/splash.lss')
    print 'Using %s.' % splashlss
    if not os.path.exists(splashlss):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % splash
        sys.exit(-1)
        
    shutil.copy(splashlss,bootloaddir)
    
    updatesimg = os.path.join(trunk,'updates.img')
    print 'Using %s.' % updatesimg
    if not os.path.exists(updatesimg):
        print 'Cannot find %s. Please regenerate it again!' % updatesimg
        sys.exit(-1)
    
    # create a temp holding directory
    pdir = tempfile.mkdtemp(prefix='primitive-')
    atexit.register(cleanup,pdir)
    
    print 'Merging root img with updates.img..'
    cmd = 'sleep 2 && mount -o loop,ro %s %s' % (updatesimg, pdir)
    mountP = subprocess.Popen(cmd,shell=True)
    mountP.wait()
    
    sys.path.append(os.path.abspath(os.path.join(pdir,'opt/kusu/lib/python')))
    sys.path.append(os.path.abspath(os.path.join(pdir,'opt/kusu/lib64/python')))
    sys.path.append(os.path.abspath(os.path.join(pdir,'opt/primitive/lib/python2.4/site-packages')))
    sys.path.append(os.path.abspath(os.path.join(pdir,'opt/primitive/lib64/python2.4/site-packages')))
    sys.path.append(os.path.abspath(os.path.join(pdir,'opt/primitive/libexec')))
    
    from path import path
    from primitive.repo.yast import YastRepo
    from primitive.fetchtool.commands import FetchCommand
    
    y = YastRepo(rootdir)
    y.handleUpdates('file://' + str(os.path.abspath(updatesimg)))
    
    cmd = 'umount %s' % pdir
    umountP = subprocess.Popen(cmd,shell=True)
    umountP.wait()

def mkISOLinuxDir(rootdir, buildprofile):
    """Creates a isolinux directory in rootdir with prefilled settings for the distro"""
    if buildprofile.kitsource:
        if not os.path.exists(buildprofile.kitsource):
            print '%s directory does not exist!' % buildprofile.kitsource
            sys.exit(-1)
        # look for the custom vmlinuz/initrd    
        vmlinuz = os.path.join(buildprofile.kitsource,'isolinux/vmlinuz')
        if not os.path.exists(vmlinuz):
            print 'No vmlinuz found in %s/isolinux !' % buildprofile.instsource
            sys.exit(-1)
        print 'Using %s as the kernel image..' % vmlinuz

        initrd = os.path.join(buildprofile.kitsource,'isolinux/initrd.img')        
        if not os.path.exists(initrd):
            print 'No initrd.img found in %s/isolinux !' % buildprofile.instsource
            sys.exit(-1)
        print 'Using %s as the initrd image..' % initrd
            
    else:
        # look for the distro's vmlinuz/initrd
        vmlinuz = os.path.join(buildprofile.instsource,'isolinux/vmlinuz')
        if not os.path.exists(vmlinuz):
            print 'No vmlinuz found in %s/isolinux !' % buildprofile.instsource
            sys.exit(-1)
        print 'Using %s as the kernel image..' % vmlinuz

        initrd = os.path.join(buildprofile.instsource,'isolinux/initrd.img')        
        if not os.path.exists(initrd):
            print 'No initrd.img found in %s/isolinux !' % buildprofile.instsource
            sys.exit(-1)
        print 'Using %s as the initrd image..' % initrd
        
        
    isolinuxbin = os.path.join(buildprofile.instsource,'isolinux/isolinux.bin')
    if not os.path.exists(os.path.join(isolinuxbin)):
        print 'No isolinux.bin found in %s/isolinux !' % buildprofile.instsource
        sys.exit(-1)
        
    isolinuxdir = os.path.join(rootdir,'isolinux')
    os.mkdir(isolinuxdir)
    shutil.copy(vmlinuz,isolinuxdir)
    shutil.copy(initrd,isolinuxdir)
    shutil.copy(isolinuxbin,isolinuxdir)

    
    # get the svn trunk layout
    trunk = getKusuTrunk()
    
    # grab the isolinux.cfg
    if buildprofile.kitsource and os.path.exists(os.path.join(buildprofile.kitsource,'isolinux/isolinux.cfg')):
        isolinuxcfg = os.path.join(buildprofile.kitsource,'isolinux/isolinux.cfg')
    else:        
        isolinuxcfg = os.path.join(trunk, \
                        'src/dists/%(distro)s/%(version)s/%(arch)s/initrd/isolinux.cfg' %
                        {'distro':buildprofile.distro,
                        'version':buildprofile.version,
                        'arch':buildprofile.arch})                    
    if not os.path.exists(isolinuxcfg):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % isolinuxcfg
        sys.exit(-1)

    print 'Using %s as the isolinux.cfg file..' % isolinuxcfg
    
    shutil.copy(isolinuxcfg,isolinuxdir)
        
    # grab the ks.cfg
    kscfg = os.path.join(trunk, \
                    'src/dists/%(distro)s/%(version)s/%(arch)s/initrd/ks.cfg' %
                    {'distro':buildprofile.distro,
                    'version':buildprofile.version,
                    'arch':buildprofile.arch})

    if not os.path.exists(kscfg):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % kscfg
        sys.exit(-1)

    shutil.copy(kscfg,rootdir)
        
    # grab the boot.msg
    bootmsg = os.path.join(trunk, \
                    'src/dists/%(distro)s/%(version)s/%(arch)s/boot.msg' %
                    {'distro':buildprofile.distro,
                    'version':buildprofile.version,
                    'arch':buildprofile.arch})
    print 'Using %s.' % bootmsg
    if not os.path.exists(bootmsg):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % bootmsg
        sys.exit(-1)
    
    shutil.copy(bootmsg,isolinuxdir)
    
    # grab the splash.lss
    splashlss = os.path.join(trunk, \
                    'src/dists/common/splash.lss')

    if not os.path.exists(isolinuxcfg):
        print 'Cannot find %s! Please ensure that the SVN source is correct!' % splash
        sys.exit(-1)
        
    shutil.copy(splashlss,isolinuxdir)
    

    # patch the initrd to include our ks.cfg
    # grab the custom ks.cfg
    customkscfg = os.path.join(trunk, \
                    'src/dists/%(distro)s/%(version)s/%(arch)s/initrd/ks.cfg' %
                    {'distro':buildprofile.distro,
                    'version':buildprofile.version,
                    'arch':buildprofile.arch})
    localinitrd = os.path.join(isolinuxdir,'initrd.img')
    print 'Patching initrd with custom ks.cfg..'
    patchInitrd(localinitrd,customkscfg)

    
def mkImagesDir(rootdir, buildprofile):
    """Creates an images directory in rootdir and place the patchfile in it"""

    # create the images directory
    imagesdir = os.path.join(rootdir, 'images')
    os.mkdir(imagesdir)
    
    # locate the updates.img file
    trunk = getKusuTrunk()
    updatesimg = os.path.join(trunk,'updates.img')
    if not os.path.exists(updatesimg):
        print 'Cannot find %s. Please regenerate it again!' % updatesimg
        sys.exit(-1)
    
    shutil.copy(updatesimg,imagesdir)

    # locate the stage2.img from the distro installation source
    stage2img = os.path.join(buildprofile.instsource,'images/stage2.img')
    if not os.path.exists(stage2img):
        print 'Cannot find %s. Please ensure that the distro installation source is valid!' % stage2img
        sys.exit(-1)

    shutil.copy(stage2img,imagesdir)
    
def mkMetaISO(rootdir, buildprofile):
    """Copies all kits found in the trunk/iso directory into the meta kit iso and generate iso"""
    
    trunk = getKusuTrunk()
    
    isodirname = os.path.abspath(os.path.join(trunk, 'iso'))
    
    print 'Looking for kits in %(dirname)s..' % {'dirname':isodirname}
    
    li = glob.glob(os.path.abspath(os.path.join(isodirname,'*.iso')))
        
    # handle the isos
    tmpdir = tempfile.mkdtemp(prefix='emerald-')
    atexit.register(cleanup,tmpdir)
    mntlist = []
    kdirs = []
    for l in li:
        namebase, ext = os.path.splitext(os.path.basename(l))
        mntpnt = os.path.join(tmpdir,namebase)
        os.mkdir(mntpnt)
        mntlist.append(mntpnt)
        cmd = 'sleep 2 && mount -o loop,ro %s %s' % (l,mntpnt)
        mountP = subprocess.Popen(cmd,shell=True)
        mountP.wait()
        kdirs.extend([os.path.join(mntpnt,d) for d in os.listdir(mntpnt) if os.path.isdir(os.path.join(mntpnt,d))])
        
    if not kdirs:
        # no kits found. cleanup and exit
        for mntpnt in mntlist:
            cmd = 'umount %s' % mntpnt
            umountP = subprocess.Popen(cmd,shell=True,cwd=tmpdir)
            umountP.wait()
            
        print 'No kits found in %s!' % isodirname
        sys.exit(-1)
        
    for kdir in kdirs:
        print 'Found %(kit)s. Including it..' % {'kit':os.path.basename(kdir)}

    # create the symlinks
    cmd = 'ln -sf %s .'
    for kdir in kdirs:
        symlinkP = subprocess.Popen(cmd % kdir,shell=True,cwd=rootdir)
        symlinkP.wait()
        
    # make the iso
    kdirs.sort()
    _isofile = '%s.iso' % '+'.join([os.path.basename(kdir) for kdir in kdirs])
    isofile = os.path.abspath(_isofile)
    
    if buildprofile.distro == 'sles' or buildprofile.distro == 'opensuse':
        makeSuSEBootISO(rootdir,isofile,buildprofile.arch)
    else:
        makeBootISO(rootdir,isofile)
    print "Generated meta kit:", isofile
    
    for mntpnt in mntlist:
        cmd = 'umount %s' % mntpnt
        umountP = subprocess.Popen(cmd,shell=True,cwd=tmpdir)
        umountP.wait()
    
    
def makeBootISO(rootdir, isofile, volname="OCS5"):
    """Create a bootable ISO isoname out of the directory rootdir. The volname can be used to specify
       the ISO label name.
    """

    if not os.path.exists(rootdir):
        print '%s not found!' % rootdir

    # nuke existing iso file if any
    if os.path.exists(isofile): os.remove(isofile)

    options = '-c isolinux/boot.cat \
    -b isolinux/isolinux.bin \
    -V %s\
    -quiet \
    -no-emul-boot \
    -boot-load-size 4 \
    -boot-info-table \
    -R -r -T -f -l \
    -o %s %s' %  (volname,isofile,rootdir)

    isoP = subprocess.Popen('mkisofs %s > /dev/null 2>&1' % options,shell=True,cwd=rootdir)
    isoP.wait()

def makeSuSEBootISO(rootdir, isofile, arch, volname="kusu"):
    """Create a bootable ISO isoname out of the directory rootdir. The volname can be used to specify
       the ISO label name.
    """

    if not os.path.exists(rootdir):
        print '%s not found!' % rootdir

    # nuke existing iso file if any
    if os.path.exists(isofile): os.remove(isofile)

    options = '-c boot/%s/loader/boot.cat \
    -b boot/%s/loader/isolinux.bin \
    -V %s\
    -quiet \
    -no-emul-boot \
    -boot-load-size 4 \
    -boot-info-table \
    -graft-points \
    -iso-level 4 \
    -pad \
    -allow-leading-dots \
    -R -r -T -f -J -l \
    -o %s %s' %  (arch, arch, volname,isofile,rootdir)

    isoP = subprocess.Popen('mkisofs %s > /dev/null 2>&1' % options,shell=True,cwd=rootdir)
    isoP.wait()

if __name__ == '__main__':
    parser = optparse.OptionParser()
    parser.add_option('-d', '--distro', dest='distro', help="Specify distro")
    parser.add_option('-v', '--version', dest='version', help="Specify version")
    parser.add_option('-a', '--arch', dest='arch', help="Specify arch")
    parser.add_option('-s', '--instsource', dest='source', help="Specify distro installation source path")
    parser.add_option('-k', '--kitsource', dest='kitsource', help="Specify kit source path with custom kernel/initrd")
    options, args = parser.parse_args()
    
    if not options.distro: 
        print 'Please specify the distro!'
        sys.exit(-1)
        
    if not options.version: 
        print 'Please specify the version!'
        sys.exit(-1)
            
    if not options.arch:
        print 'Please specify the arch!'
        sys.exit(-1)
        
    if not options.source: 
        print 'Please specify the distro installation source path!'
        sys.exit(-1)
    
    # create the buildprofile object
    buildprofile = BuildProfile()
    buildprofile.instsource = options.source
    buildprofile.distro = options.distro
    buildprofile.version = options.version
    buildprofile.arch = options.arch
    
    if options.kitsource:
        buildprofile.kitsource = options.kitsource
    
    # create scratch directory
    tdir = tempfile.mkdtemp(prefix='emerald-')
    atexit.register(cleanup,tdir)

    if buildprofile.distro == 'sles' or buildprofile.distro == 'opensuse':
        makeSuseBootDir(tdir, buildprofile)
    else:
        mkISOLinuxDir(tdir, buildprofile)
        mkImagesDir(tdir, buildprofile)
        
        # stamp .discinfo
        discinfo = os.path.join(tdir,'.discinfo')
        f = open(discinfo,'w')
        f.close()
    
    mkMetaISO(tdir, buildprofile)
    sys.exit(0)