#coding:utf-8
import os
import re
import sys
import time
import string
import socket
import shutil
import sunhpc
import pexpect
import tempfile
import subprocess

class Builder:
    def __init__(self, cmd):
        self.cmd     = cmd
        self.config  = None
        self.tempdir = os.getcwd()

    def mktemp(self):
        return tempfile.mktemp(dir=self.tempdir)
        
    def mkisofs(self, isoName, rollName, diskName, rollDir):
        self.cmd.msg('Building ISO image for %s ...' % diskName)

        if self.config.isBootable():
            extraflags = self.config.getISOFlags()
        else:
            extraflags = ''

        volname = '%s %s' % (rollName, diskName)
        if len(volname) > 32:
            volname = volname[0:32]
            
        cwd = os.getcwd()
        cmd = 'mkisofs -V "%s" %s -input-charset utf-8 -max-iso9660-filenames -r -T -f -o %s .' % \
            (volname, extraflags, os.path.join(cwd, isoName))

        os.chdir(rollDir)
        self.cmd.system(cmd, std=0)
        os.chdir(cwd)

    def discinfo(self, dir, name, arch, id=1):
        # .discinfo 格式
        # 第一行是时间戳
        # 第二行是碟片名称
        # 第三行是架构信息
        # 第四行是碟片序号 1,2,3 连续三张碟片.
        disc = os.path.join(dir, '.discinfo')
        if os.path.isfile(disc):
            os.unlink(disc)

        with open(disc, 'w') as f:
            f.write('%f\n' % time.time())
            f.write('%s\n' % name)
            f.write('%s\n' % arch)
            f.write('%d\n' % id)

class RollBuilder(Builder, sunhpc.core.dist.Arch):

    def __init__(self, cmd, xmlfile, sign, kernel):
        Builder.__init__(self, cmd)
        sunhpc.core.dist.Arch.__init__(self)

        self.config = sunhpc.core.files.RollInfoFile(xmlfile)
        self.setArch(self.config.getRollArch())

        self.cmd    = cmd
        self.sign   = sign   # 是否需要签名.
        self.kernel = kernel # 提供的Linux iso挂载的根目录.

    def mkisofs(self, isoName, rollName, diskName):
        Builder.mkisofs(self, isoName, rollName, diskName, diskName)

    def getRPMS(self, path):

        dict = {}
        tree = sunhpc.core.files.Tree(os.path.join(os.getcwd(), path))
        for dir in tree.getDirs():
            for file in tree.getFiles(dir):
                try:
                    file.getPackageName()
                except AttributeError:
                    continue # skip all non-rpm files
                    
                # Skip RPMS for other architecures
                if file.getPackageArch() not in self.getCPUs():
                    continue
                    
                # Resolve package versions
                name = file.getUniqueName()
                if not name in dict or file >= dict[name]:
                    dict[name] = file
                    
        list = []
        for e in dict.keys():
            list.append(dict[e])
        return list

    def signRPM(self, rpm):

        cmd = "rpm -q --qf '%%{BUILDHOST}' -p %s" % rpm.getFullName()
        buildhost = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).stdout.readline()
        hostname  = socket.gethostname()

        if buildhost == hostname:
            print (rpm.getFullName())
            print ('Need to resign for the rpm files')

    def treeinfo(self, name):

        with open(os.path.join(name, '.treeinfo'), 'w') as f:
            ti = '[general]\n'
            ti += 'name = Sunhpc-%s\n' % sunhpc.version_major
            ti += 'family = SunhpcOS\n'
            ti += 'timestamp = %.2f\n' % time.time()
            ti += 'variant = \n'
            ti += 'version = %s\n' % sunhpc.version_major
            ti += 'packagedir =\n' 
            ti += 'arch = x86_64\n\n'

            ti += '[stage2]\n'
            ti += 'mainimage = images/install.img\n\n'

            ti += '[images-x86_64]\n'
            ti += 'kernel = images/vmlinuz-%s-%s\n' % (sunhpc.version, self.cmd.arch)
            ti += 'initrd = images/initrd-%s-%s\n' % (sunhpc.version, self.cmd.arch)
            ti += 'boot.iso = images/boot.iso\n\n'
            f.write(ti)

    def makeBootOS(self, disk, root):
        self.cmd.msg('\tCreate roll comps repodata ......')
        cwd = os.getcwd()
        #
        # base: roll/kernel/version/x86_64/RedHat/base
        #
        base = os.path.join(cwd, root, 'RedHat', 'base')
        os.makedirs(base)

        #
        # 将光盘中的repodata/中的*-*-*comps.xml文件放入base目录中.
        # comps = os.path.join(cwd, disk, 'comps.xml')
        #
        comps = os.path.join(base, 'comps.xml')
        if not os.path.exists(comps):
            compfile = ''
            repodirs = os.path.join(self.kernel, 'repodata')
            for c in os.listdir(repodirs):
                if c.endswith('.xml') and 'comps' in c.split('.')[0]:
                    compfile = os.path.join(repodirs, c)
                    break

            if not os.path.exists(compfile):
                self.cmd.msg('The *-comps.xml file was not found in %s' % self.kernel, 'a')
                
            shutil.copy(compfile, comps)

        # 通过.treeinfo 中的upgrade特性创建自定义镜像文件.
        product = os.path.join(cwd, disk, 'updates')
        if not os.path.exists(product):
            os.makedirs(product)
            self.treeinfo(disk)

        images = os.path.join(cwd, disk, 'images')
        if not os.path.exists(images):
            shutil.copytree(os.path.join(cwd, 'isolinux'), os.path.join(disk, 'images'))

        isolinux = os.path.join(cwd, disk, 'isolinux')
        if not os.path.exists(isolinux):
            isodirs = os.path.join(self.kernel, 'isolinux')
            shutil.copytree(isodirs, isolinux)
        
        if not os.path.exists(comps):
            self.cmd.msg('\n\tCould not find a comps.xml file. - %s' % comps, 'e')
            sys.exit(-1)

        self.createrepo(disk, comps)

    def createrepo(self, disk, comps):
        self.cmd.msg('Creating %s repository......' % disk)
        cwd  = os.getcwd()
        base = os.path.join(cwd, disk)
        os.chdir(base)

        createrepo = '/usr/share/createrepo/genpkgmetadata.py'
        if not os.path.exists(createrepo):
            self.cmd.msg('First yum install createrepo', 'a')

        # 使用works参数,需要有足够的tmpdir空间.
        tmpdir = os.getenv("TMPDIR")
        os.putenv("TMPDIR",".")
        cmd = "%s --groupfile %s --workers 8 --quiet ." % (createrepo, comps)
        subprocess.call(cmd, shell=True)
        if tmpdir is not None:
            os.putenv("TMPDIR",tmpdir)
        else:
            os.unsetenv("TMPDIR")
        os.chdir(cwd)

    def getExternalRPMS(self):
        # 通过解析kickstart xml脚本提取出来的所有rpms.
        rpms = ['gcc']
        '''
        selected = []
        for rpm in os.popen('%s/sbin/runpy2 %s' % (
            os.environ.get('SUNHPC_HOME'), ' '.join(rpms))):
            selected.extend(eval(rpm))
        '''
        sunhpc = []
        nonsun = []
        return ([], [])

    def getReqiredRPMS(self):
        self.cmd.msg('Please wait while collecting rpm packages...', 'e')

        rpms = ['minimal', '@base', '@core', '@development']
        required = []
        for rpm in os.popen('%s/sbin/runpy2 %s' % (
            os.environ.get('SUNHPC_HOME'), ' '.join(rpms))):
            required.extend(eval(rpm))

        if os.path.exists('RPMS'):
            shutil.rmtree('RPMS')

        os.makedirs('RPMS')

        self.cmd.msg('Copying iso boot required rpm files...', 'e')
        packDirs = os.path.join(self.kernel, 'Packages')
        if not os.path.exists(packDirs):
            self.cmd.msg('Cannot find %s dirs,may not be booting!' % packDirs, 'w')

        packages = {}
        packTree = sunhpc.core.files.Tree(self.kernel)
        for dir in packTree.getDirs():
            if dir != 'Packages': continue

            for file in packTree.getFiles(dir):
                packages[file.getBaseName()] = file

        #os.symlink(file.getFullName(), os.path.join(path, file.getName()))
        for p in required:
            if p in packages:
                os.symlink(
                    packages[p].getFullName(), 
                    os.path.join('RPMS', packages[p].getName())
                )
            else:
                self.cmd.msg('The %s package could not be found.' % p)

    def run(self):

        rpmlist = []
        if self.config.hasRPMS():
            # 收集当前RPMS目录里面的rpm文件.
            rpmlist.extend(self.getRPMS('RPMS'))
            rpmlist.extend(self.getRPMS('PXES'))

        if self.config.hasSRPMS():
            # 收集当前SRPMS目录里面的rpm文件.
            rpmlist.extend(self.getRPMS('SRPMS'))
        
        # 出现错误: Header V3 RSA/SHA256 Signature, key ID f4a80eb5: NOKEY
        # 需要导入gpg-key
        # rpm --import /etc/pki/rpm-gpg/RPM-GPG-KEY-CentOS-7
        if rpmlist:
            cmd = "rpm -q --qf '%%{BUILDHOST}' -p %s" % rpmlist[0].getFullName()
            fix = "rpm --import /etc/pki/rpm-gpg/RPM-*"
            ret = self.cmd.shcmd(cmd, code=True)
            if ret['c'] != 0: 
                self.cmd.shcmd(fix, code=True)

        # 进行验签,如果有我们自己制作的rpm文件,需要进行重签
        # 进行了签名以后,在yum的配置时候就可以开启验签了.
        if self.sign:
            self.cmd.msg('Starting the rpm sign...')
            for rpm in rpmlist:
                self.signRPM(rpm)

        required = []
        if self.config.hasRolls():
            self.cmd.msg('Starting making the rolls iso...')
            (required, optional) = self.getExternalRPMS()
            for file in rpmlist:
                required.append(file)

            # required is total packages.
            self.cmd.msg('\tRequired Packages %s' % len(required))
            self.cmd.msg('\tOptional Packages %s' % len(optional))
            for file in required:
                rpmlist.append(file)
            rpmlist.extend(optional)

        # 执行完成getExternalRPMS(),required不应该为0
        if len(required) == 0:
            self.cmd.msg('This disk is optional (extra rpms)', 'w')

        # ISO Roll基础目录.
        name = 'disk'
        root = os.path.join(name, self.config.getRollName(),
                self.config.getRollVersion(), self.config.getRollArch())

        os.makedirs(root)
        os.makedirs(os.path.join(root, 'RedHat', 'RPMS'))
        os.makedirs(os.path.join(root, 'SRPMS'))

        # Symlink in all the RPMS and SRPMS
        for file in rpmlist:
            try:
                arch = file.getPackageArch()
            except:
                continue

            if arch == 'src':
                file.symlink(os.path.join(root, 'SRPMS', file.getName()))
            else:
                file.symlink(os.path.join(root, 'RedHat', 'RPMS', file.getName()))

            if file in required:
                del required[required.index(file)]

        # Copy the Roll XML file onto all the disks
        shutil.copy(self.config.getFullName(), root)
            
        # Create the .discinfo file
        self.discinfo(name, self.config.getRollName(), self.config.getRollArch())

        isoname = '%s-%s-%s.%s.%s.iso' % (
                self.config.getRollName(),
                self.config.getRollVersion(),
                self.config.getRollRelease(),
                self.config.getRollArch(),
                name)
                
        if self.config.isBootable() == 1:
            self.cmd.msg('Configuring %s kickstart bootable ...' % self.config.getRollName())
            self.makeBootOS(name, root)
            #self.makeBootGrub(self.config.getRollName(), name)
            
        self.mkisofs(isoname, self.config.getRollName(), name)
        self.cmd.msg("%s create finished." % isoname)

    def makeBootGrub(self, rollname, name):
        # CentOS7 以后的grub引导菜单里面的Label需要转义空格.
        config  = os.path.join('kernel', 'isolinux', 'isolinux.cfg')
        srckey  = 'inst.stage2=hd:LABEL=CentOS\\x207\\x20x86_64'
        dstkey  = 'inst.stage2=hd:LABEL=%s\\x20%s ' % (rollname, name)
        dstkey += 'net.ifnames=0 biosdevname=0'

        text = self.cmd.replace(config, srckey, dstkey)
        with open(os.path.join(name, 
                'isolinux', 'isolinux.cfg'), 'w') as f:
            for i in text:
                f.write(i)

    def makeBoot(self):
        if not self.config.isBootable():
            return

        file_list = ['initrd.img', 'vmlinuz', 'squashfs.img', 'efiboot.img']
        for f in file_list:
            require = self.cmd.shcmd('find %s -type f -name %s' % (self.kernel, f))['o'].strip()
            if not require:
                self.cmd.msg('.e.g sunhpc create roll roll-*.xml boot=/mnt/cdrom', 'e')
                self.cmd.msg('must provide an source linux cd path build to boot. %s does not exist.' % self.kernel, 'a')

        self.cmd.msg('Configuring %s pxelinux bootable ...' % self.config.getRollName())
        if not os.path.exists('isolinux'):
            os.makedirs('isolinux')

        pxelist = ['product', 'updates', 'efiboot', 'vminitrd', 'squashfs', 'kickstart', 'client']
        for i in pxelist:
            self.cmd.command('create.pxelinux.%s' % i)
            src = 'source/images'
            if not os.path.exists(src):
                continue
            for s in os.listdir(src):
                shutil.copyfile(os.path.join(src, s), os.path.join('isolinux', s))

class Command(sunhpc.commands.create.command):
    """
    创建新的Roll.
    <arg type='file' name="*.xml">
    Either a list of Roll ISO files or the name of a single Roll XML
    description file.  If a list of Roll ISO files to be merge together 
    into a single Roll.  Otherwise the single argument is assumed to
    be the name of the XML file generated by the top level Makefile in
    the Roll's source.
    </arg>

    <param type="bool" name="sign">
    Whether to sign the rpms
    </param>

    <param type="Path" name="boot">
    Provides an ISO mount directory for the original system distribution.
    </param>

    <example cmd='create roll roll-base.xml'>
    Create an new roll iso file
    </example>

    <example cmd='create roll roll-CentOS.xml boot=/mnt/cdrom'>
    Create an new boot roll iso file.
    </example>

    <related>create xml name=CentOS version=7.9.2009 boot=1'</related>
    """

    def run(self, param, args):

        (kernel, sign) = self.fillParams([('boot', ''), ('sign', 'no')])

        sign = self.str2bool(sign)
        if len(args) != 1:
            self.msg('must supply roll xml config file. or use command create it.', 'e')
            self.msg('\t# sunhpc create xml name=CentOS version=7.9.2009 boot=1', 'a')

        base, ext = os.path.splitext(args[0])
        if ext != '.xml':
            self.msg('missing xml file - %s' % args[0], 'a')
            
        cwd     = os.getcwd()
        builder = RollBuilder(self, args[0], sign, kernel)
                
        # create boot iso symlink
        rpms = os.path.join(kernel, 'Packages')
        if not os.path.islink('RPMS'):
            os.symlink(rpms, 'RPMS')

        builder.makeBoot()
        if os.path.exists('pxeboot') and not os.path.islink('PXES'):
            os.symlink('pxeboot', 'PXES')

        builder.run()
        self.msg('Use following command init pxe:', 'i')
        self.msg('    sunhpc pxelinux build xxx.x86_64.disk.iso', 'i')




