#coding:utf-8
import os
import re
import sys
import xml
import time
import sunhpc
import shutil
import socket
import tempfile
import subprocess 
class BuildError(Exception):
    pass

class Builder:
    
    def __init__(self):
        self.verbose = 0
        self.debug   = 0

    def build(self):
        pass

    def setVerbose(self, level=1):
        self.verbose = level

    def setDebug(self, level=1):
        self.debug = level

class MirrorBuilder(Builder):

    def __init__(self, m):
        Builder.__init__(self)
        self.mirrors    = m

    def build(self):
        for m in self.mirrors:
            dirs = []
            if m.getRemoteReleasePath():
                dirs.append(m.getRemoteReleasePath())
            for dir in dirs:
                self.buildMirror(m.getHost(), dir)


    def buildMirror(self, host, path):
        # Try FTP first, failover to HTTP
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            sock.connect((host, 21))
            sock.close()
            cmd = 'wget -m -nv ftp://%s//%s/' % (host, path)
        except socket.error:
            cmd = 'wget -m -nv -np http://%s//%s/' % (host, path)
            sock = None

        if self.verbose or self.debug:
            print (cmd)
        if not self.debug:
            subprocess.call(cmd, shell=True)

class DistributionBuilder(Builder):

    def __init__(self, dist, links=1):
        Builder.__init__(self)
        self.cmd            = None
        self.dist           = dist
        self.quiet          = None
        self.useLinks       = links
        self.compsPath      = None
        self.useRolls       = {}
        self.allRolls       = 1
        self.onlyRolls      = 0
        self.version        = '1.0'
        self.calcmd5        = 1
        self.withSiteProfiles   = 0
        for mirror in self.dist.getMirrors():
            if not mirror.isBuilt():
                mirror.build()

        if not self.dist.isBuilt():
            self.dist.build()

    def setCommand(self, command):
        self.cmd = command
    def setQuiet(self, quiet):
        self.quiet = quiet

    def setRolls(self, list, only=0):
        if list:
            for e in list:
                self.useRolls[e[0]] = (e[1], e[2])
            self.allRolls = 0
        else:
            self.useRolls = {}
            self.allRolls = 1
        self.onlyRolls = only

    def setVersion(self, ver):
        self.version = ver

    def setSiteProfiles(self, bool):
        self.withSiteProfiles = bool
        
    def setCalcMD5(self, bool):
        self.calcmd5 = bool
        
    def clean(self):
        self.dist.getTree('release').apply(self.cleaner)

    def useRoll(self, key, ver, arch):
        "Returns true if we should include this roll"
        if arch == self.dist.arch:
            if self.allRolls:
                return 1
        if key in self.useRolls:
            version, enabled = self.useRolls[key]
            if enabled and version == ver:
                return 1
        return 0

    def getRollBaseFiles(self):
        # 这里收集rolls中文件
        # roll 名字,架构,版本,都必须和数据库匹配才可以收集到信息.
        files = []
        for m in self.dist.getMirrors():
            # m: dist-Mirror class
            for key, value in m.getRolls().items():
                # key-value : base   [('x86_64', '7.0.0')]
                # key-value : kernel [('x86_64', '7.0.0')]
                for arch, ver in value:
                    if self.useRoll(key, ver, arch):
                        if not self.quiet:
                            self.cmd.msg('    including roll (%s,%s) "%s" ...' % (ver, arch, key))
                        files.extend(m.getRollBaseFiles(key, ver, arch))
        return files
        
    def getRollRPMS(self):
        rpms = []
        for m in self.dist.getMirrors():
            for key, value in m.getRolls().items():
                for arch, ver in value:
                    if self.useRoll(key, ver, arch):
                        if not self.quiet:
                            self.cmd.msg('    including roll (%s,%s) "%s" ...' % (ver, arch, key))
                        rpms.extend(m.getRollRPMS(key, ver, arch))
        return rpms
        
    def getRollSRPMS(self):
        rpms = []
        for m in self.dist.getMirrors():
            for key, value in m.getRolls().items():
                for arch, ver in value:
                    if self.useRoll(key,ver,arch):
                        if not self.quiet:
                            self.cmd.msg('    including roll (%s,%s) "%s" ...' % (ver, arch, key))
                        rpms.extend(m.getRollSRPMS(key, ver, arch))
        return rpms
        
    def buildRPMSList(self):
        rpms = self.getRollRPMS()
        for mirror in self.dist.getMirrors():
            rpms.extend(mirror.getRPMS())
        if not self.onlyRolls:
            rpms.extend(self.dist.getContribRPMS())
            rpms.extend(self.dist.getLocalRPMS())
        if not os.path.isdir(self.dist.getForceRPMSPath()):
            os.makedirs(self.dist.getForceRPMSPath())
        else:
            rpms.extend(self.dist.getForceRPMS()) 
        return rpms

    def buildSRPMSList(self):
        rpms = self.getRollSRPMS()
        for mirror in self.dist.getMirrors():
            rpms.extend(mirror.getSRPMS())
        rpms.extend(self.dist.getContribSRPMS())
        rpms.extend(self.dist.getLocalSRPMS())
        return rpms

    def buildRollLinks(self):
        """Links all rolls from our mirrors into rocks-dist/rolls/"""
    
        print("Building Roll Links")
        rollLocation = self.dist.getRollsPath()
        subprocess.call('mkdir -p %s' % rollLocation, shell=True)

        rolls = []
        for mirror in self.dist.getMirrors():
            rolldir = mirror.getRollsPath()
            if not os.path.exists(rolldir):
                continue
            for d in os.listdir(rolldir):
                rollpath = os.path.join(rolldir,d)
                if os.path.isdir(rollpath):
                    rolls.append(rollpath)

        here = os.getcwd()
        os.chdir(rollLocation)
        for r in rolls:
            subprocess.call('ln -sf %s .' % (r), shell=True)
        os.chdir(here)
    
    def buildBase(self):
        if not self.quiet:
            baselen = len(self.resolveVersions(self.getRollBaseFiles()))
            self.cmd.msg('Resolving versions (base - %s)' % baselen)

        self.dist.setBaseFiles(self.resolveVersions(self.getRollBaseFiles()))

    def touchCriticalFiles(self, m, key, ver, arch):
        criticalfiles = [ 'anaconda', 'anaconda-runtime', 'kudzu', 'kudzu-devel' ]
        for rpm in m.getRollRPMS(key,ver,arch):
            try:
                if rpm.getPackageName() in criticalfiles:
                    rpm.timestamp = int(time.time())
            except:
                pass

    def includeCriticalRPMS(self):
        if not self.quiet:
            self.cmd.msg('    including "critical" RPMS')
        for m in self.dist.getMirrors():
            for key, value in m.getRolls().items():
                if key != 'base':
                    continue
                for arch, ver in value:
                    if self.useRoll(key, ver, arch):
                        self.touchCriticalFiles(m,key,ver,arch)

    def buildRPMS(self):
        if not self.quiet:
            rpms = len(self.resolveVersions(self.buildRPMSList()))
            self.cmd.msg('Resolving versions (RPMs  - %s)' % rpms)
        self.dist.setRPMS(self.resolveVersions(self.buildRPMSList()))

    def buildSRPMS(self):
        if not self.quiet:
            srpms = len(self.resolveVersions(self.buildSRPMSList()))
            self.cmd.msg('Resolving versions (SRPMs - %s)' % srpms)
        self.dist.setSRPMS(self.resolveVersions(self.buildSRPMSList()))

    def insertNetstage(self):
        if not self.quiet:
            self.cmd.msg('Rebuiling the images...')
        
        # cmd = 'rm -f %s/RedHat/base/install.img' % (self.dist.getReleasePath())
        # subprocess.call(cmd, shell=True)

        imgFiles = ['install', 'initrds', 'efiboot', 'product', 'updates']
        for i in imgFiles:
            if not self.quiet:
                self.cmd.msg('    Applying %s images...' % i)
            try:
                self.applyRPM('sunhpc-anaconda-%s' % i, self.dist.getReleasePath())
            except:
                self.cmd.msg("Couldn't find the package sunhpc-anaconda-%s" % i, 'w')

    def build(self):
        #Start SUNHPC Build
        self.cmd.msg('Staring build the sunhpc distro release ...')

        self.clean()
        self.dist.syncMirror()

        # 开始收集相关文件.
        self.buildBase()
        self.includeCriticalRPMS()
        self.buildRPMS()
        self.buildSRPMS()

        if not self.quiet:
            self.cmd.msg('Start creating files ...')
        if self.useLinks:
            if not self.quiet:
                self.cmd.msg('    (symbolic links - fast)')
        else:
            if not self.quiet:
                self.cmd.msg('    (use deep copy  - slow)')

        self.dist.getReleaseTree().apply(self.builder)
        self.dist.getReleaseTree().apply(self.normalizer)

        # install.img and updates.img
        self.insertNetstage()

        # install roll-*-kickstart
        self.buildKickstart()

        # write treeinfo
        self.buildTreeInfo()

        if not self.quiet:
            self.cmd.msg('Calling Yum genpkgmetadata.py')
        self.createrepo()

        if not self.quiet:
            self.cmd.msg('Rebuilding Product Image including md5 sums')
        #self.buildProductImg()

        if not self.quiet:
            self.cmd.msg('Creating Directory Listing')
        self.makeDirListing()

        if not self.quiet:
            self.cmd.msg('The sunhpc cluster system distribution build is complete.')

    def buildTreeInfo(self):
        if not self.quiet:
            self.cmd.msg('    installing profiles "treeinfo" ...')
        treeinfo = os.path.join(self.dist.getReleasePath(), '.treeinfo')
        self.cmd.makeTreeInfo(treeinfo)

    def buildKickstart(self):
        if not self.quiet:
            self.cmd.msg('Installing XML Kickstart profiles')

        build   = self.dist.getBuildPath()
        for rpm in self.dist.getRPMS():
            tok = rpm.getBaseName().split('-')
            if tok[0] != 'roll':
                continue
            try:
                k = tok.index('kickstart')
                rollname = '-'.join(tok[1:k])
            except ValueError:
                continue
                
            if not self.quiet:
                self.cmd.msg('    installing profiles from "%s" ...' % rollname)
            self.applyRPM(rpm.getBaseName(), build)

        # install the cgi-module-kickstart
        #self.applyRPM('cgi-module-kickstart', build)

        # Copy local profiles into the distribution.
        if self.withSiteProfiles:
            if not self.quiet:
                self.cmd.msg('    installing profiles from "site" ...')
            tree = self.dist.getSiteProfilesTree()
            for dir in tree.getDirs():
                for file in tree.getFiles(dir):
                    path = os.path.join(build, dir)
                    if not os.path.isdir(path):
                        os.makedirs(path)
                    shutil.copy(file.getFullName(),
                        os.path.join(path, file.getName()))
                    # make sure apache can read site XML
                    file.chmod(0o664)

    def applyRPM(self, name, root, flags=''):
        rpm = None
        try:
            rpm = self.dist.getRPM(name)
        except sunhpc.core.dist.DistRPMList as e:
            for r in e.list:
                if r.getPackageArch() == self.dist.getArch():
                    rpm = r
                    break

        if not rpm:
            raise ValueError("could not find %s" % name)

        dbdir = os.path.join(root, 'var', 'lib', 'rpm')
        if not os.path.isdir(dbdir):
            os.makedirs(dbdir)

        reloc = subprocess.call("rpm -q --queryformat '%{prefixes}\n' -p " +
                        rpm.getFullName() + "| grep none > /dev/null", shell=True)

        cmd = 'rpm -i --ignoresize --nomd5 --force --nodeps --ignorearch '
        cmd += '--dbpath %s ' % dbdir
        if reloc:
            cmd = cmd + '--prefix %s %s %s' % (root, flags, rpm.getFullName())
        else:
            cmd = cmd + '--badreloc --relocate /=%s %s %s' % (root, flags, rpm.getFullName())
        retval = subprocess.call(cmd + ' > /dev/null 2>&1', shell=True)

        shutil.rmtree(os.path.join(root, 'var'))
        if retval == 256:
            raise BuildError("could not apply RPM %s" % (name))
        return retval

    def buildProductImg(self):
        product = '../../images/product.img'
        productfilesdir = os.path.join(self.dist.getBuildPath(), 'include')
        if not os.path.exists(productfilesdir):
            #
            # there are no 'product' files, so there's nothing to do.
            # let's just return
            #
            return

        if not self.quiet:
            self.cmd.msg('    Create product images...')
        cwd = os.getcwd()
        #
        # make an MD5 checksum for all files in the distribution
        # the 'sed' command strips off the leading "./" from the pathnames
        # don't include the build, SRPMS and force directories
        #
        os.chdir(self.dist.getReleasePath())
        if self.calcmd5:
            cmd = '/usr/bin/md5sum `find -L . -type f | sed "s/^\.\///" | '
            cmd += 'egrep -v "^build|^SRPMS|^force" | egrep -v "rpm$"` '
            cmd += '> %s/packages.md5' % (productfilesdir)
        else:
            cmd = 'touch %s/packages.md5' % (productfilesdir)

        subprocess.call(cmd, shell=True)
        #
        # create the product.img file
        #
        os.chdir(productfilesdir)

        if not os.path.exists('../../images'):
            os.makedirs('../../images')

        subprocess.call('rm -f %s' % (product), shell=True)
        cmd = '/sbin/mksquashfs packages.md5 applets pyanaconda '
        cmd += '%s ' % (product)
        cmd += '-keep-as-directory > /dev/null 2>&1'
        subprocess.call(cmd,shell=True)

        if os.path.exists(product):
            #
            # on a server installation (e.g., frontend), mksquashfs
            # fails, but it is not important that product.img is built
            # during the installation. product.img was already downloaded
            # off the CD, so it will not be needed for the remainder of
            # the server installation.
            #
            os.chmod(product, 0o664)
        os.chdir(cwd)
        return

    def createrepo(self):
        if not self.quiet:
            self.cmd.msg('Creating repository ......')

        cwd = os.getcwd()
        releasedir = self.dist.getReleasePath()
        os.chdir(releasedir)
        #
        # first check in the install environment (/tmp/updates), then
        # look in the 'normal' place (on a running frontend).
        #
        createrepo = '/tmp/updates/usr/share/createrepo/genpkgmetadata.py'
        if not os.path.exists(createrepo):
            createrepo = '/usr/share/createrepo/genpkgmetadata.py'
        if not self.quiet:
            self.cmd.msg('    Using the genpkgmetadata create ...')
        groupfile = "%s/RedHat/base/comps.xml" % releasedir
        if os.path.exists(groupfile):
            gf = "--groupfile %s/RedHat/base/comps.xml " % (releasedir)
        else:
            self.cmd.msg("Couldn't find the groupfile %s" % groupfile, 'w')
            self.cmd.msg("\tIf you are bootstrapping, this is not a problem", 'w')
            gf = " "

        tmpdir = os.getenv("TMPDIR")

        # worker.py (Called by genpkgmetadata) needs tmp space
        os.putenv("TMPDIR",".")
        subprocess.call('%s ' % (createrepo) + gf + ' --workers 8 ' + '--quiet .', shell=True)

        if tmpdir is not None:
            os.putenv("TMPDIR",tmpdir)
        else:
            os.unsetenv("TMPDIR")
        os.chdir(cwd)
        return

    def makeDirListing(self):   
        #
        # 只有在/works/sunhpc/install 这个目录
        # 执行sunhpc create distro 的时候才会执行此处函数.
        #
        path = os.path.join(self.dist.getRootPath(), 'rolls')
        if os.path.exists(path):
            filename = os.path.join(path, 'index.cgi')

            file = open(filename, 'w')
            file.write('%s' % (directory_listing_cgi))
            file.close()

            os.chmod(path, 0o755)
            os.chmod(filename, 0o755)
        return

    def cleaner(self, path, file, root):
        if not root:
            root = self.dist.getReleasePath()
        dir = os.path.join(root, path)
        if dir not in [ self.dist.getForceRPMSPath() ]:
            os.unlink(os.path.join(dir, file.getName()))

    def builder(self, path, file, root):
        # path : 'RedHat/RPMS'
        if not root:
            # root : /root/test/tmpxxxxx/x86_64
            root = self.dist.getReleasePath()
        dir  = os.path.join(root, path)
        fullname = os.path.join(dir, file.getName())

        if file.getFullName() == fullname:
            return

        if not os.path.isdir(dir):
            os.makedirs(dir)

        if self.useLinks:
            file.symlink(fullname, self.dist.getRootPath())
        else:
            if os.path.islink(file.getFullName()):
                os.symlink(os.readlink(file.getFullName()), fullname)
            else:
                shutil.copy(file.getFullName(), fullname)
                os.utime(fullname, (file.getTimestamp(), file.getTimestamp()))

    def normalizer(self, path, file, root):
        if not root:
            root = self.dist.getReleasePath()
        dir  = os.path.join(root, path)
        fullname = os.path.join(dir, file.getName())
        if file.getFullName() != fullname:
            file.setFile(fullname)

    def resolveVersions(self, files):
        dict = {}
        for e in files:
            name = e.getUniqueName() # name w/ arch string appended
            if name not in dict or e >= dict[name]:
                dict[name] = e

        list = []
        for e in dict.keys():
            list.append(dict[e])
        return list

    def setComps(self, path):
        self.compsPath = path

directory_listing_cgi = """#!/opt/sunpy3/bin/python3
import os
try:
    dir = os.environ['DOCUMENT_ROOT'] + os.environ['REQUEST_URI']
except:
    dir = '.'

if os.path.isfile(dir):
    dir, filename = os.path.split(dir)

out = ''
out += '<html>'
out += '<body>'
out += '<table>'

#
# test environ value
#for i in os.environ:
#    out += '\n%s --- %s\n' % (i, os.environ[i])

listing = os.listdir(dir)
listing.sort(key=str.lower)
for file in listing:
    if file not in [ 'index.cgi' ]:
        out += '<tr><td>\\n'

        if os.path.isdir(os.path.join(dir, file)):
            out += '<a href="%s/">%s/</a>\\n' % (file, file)
        else:
            out += '<a href="%s">%s</a>\\n' % (file, file)

        out += '</td></tr>'
        out += '\\n'

out += '</table>'
out += '</body>'
out += '</html>'

print ('Content-type: text/html')
print ('Content-length: %d' % (len(out)))
print ('')
print (out)
"""
