#-*- coding: utf-8 -*-

# Copyright 2010 Calculate Ltd. http://www.calculate-linux.org
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

__version__ = "2.2.1"
__app__ = "calculate-builder"

import os
import re
import sys
import traceback
from os import path
from cl_utils import process,pathJoin,getRunCommands,getTupleVersion,\
                        childMounts,_toUNICODE,isMount,typeFile, \
                        removeDir
from subprocess import STDOUT,PIPE
from cl_print import color_print
from cl_datavars import DataVars
from shutil import copy2 as copy_with_perm
from cl_distr import IsoDistributive, DirectoryDistributive, \
                    DistributiveError
from cl_template import template,iniParser
from cl_vars_share import varsShare
from datetime import datetime

from cl_kernel_utils import KernelConfig,InitRamFs

from server.utils import dialogYesNo
import cl_overriding
import hashlib
from cl_lang import lang
lang().setLanguage(sys.modules[__name__])

class printNoColor:
    def colorPrint(self,attr,fg,bg,string):
        sys.stdout.write(string)

class BuilderError(Exception):
    """Installation Error"""
    builder_errors = ""

    def installExit(self,*args,**kwars):
        raise BuilderError(self.__class__.builder_errors)

    def overprintERROR(self,error):
        self.__class__.builder_errors += str(error) + "\n"

    def getBuilderErrors(self):
        return self.__class__.builder_errors

    def popBuilderErrors(self):
        res = self.__class__.builder_errors
        self.__class__.builder_errors = ""
        return res

    def getOverrideMethods(self):
        return self.installExit, self.overprintERROR

cl_overriding.exit, cl_overriding.printERROR = \
    BuilderError().getOverrideMethods()

class DataVarsBuilder(DataVars):
    """Variable class for installation"""

    def importBuilder(self, **args):
        '''Get variables for builder'''
        # section name in calculate.env
        envSection = "builder"
        # import builder variables
        self.importData(envSection, ('cl_vars_builder','cl_fill_builder'))

class cl_builder(color_print):
    """Primary class for image manipulation"""

    def __init__(self):
        self.clVars = None
        self.startMessage = ""
        self.startMessage = ""
        self.clTempl = None
        self.force = False
        self.assembleIso = False
        self.envFile = '/etc/calculate/assemble.env'

    def setNoColor(self):
        self.color = False

    def initVars(self):
        """Primary initialization of variables"""
        self.clVars = DataVarsBuilder()
        self.clVars.importBuilder()
        self.clVars.flIniFile()

    def applyTemplatesForSquash(self,directory):
        """Apply templates for root of system."""
        #self.clVars.Set("cl_root_path","/", True)
        self.clVars.Set("cl_action","squash", True)
        self.clVars.Set("cl_chroot_path",directory, True)
        self.clTempl = template(self.clVars)
        dirsFiles = self.clTempl.applyTemplates()
        self.clTempl.closeFiles()
        if self.clTempl.getError():
            raise BuilderError(self.clTempl.getError())
        else:
            return dirsFiles

    def applyTemplatesForIso(self,directory):
        """Apply templates for root of system."""
        #self.clVars.Set("cl_root_path","/", True)
        self.clVars.Set("cl_action","iso", True)
        chrootPath=pathJoin(directory,
                            self.clVars.Get('cl_builder_squash_path'))
        self.clVars.Set("cl_chroot_path",chrootPath, True)
        self.clVars.Set("cl_root_path","..", True)
        self.clTempl = template(self.clVars)
        dirsFiles = self.clTempl.applyTemplates()
        self.clTempl.closeFiles()
        if self.clTempl.getError():
            raise BuilderError(self.clTempl.getError())
        else:
            return dirsFiles

    def printMessageForTest(self,message):
        """Print waiting message and OK or Error by func result"""
        message = "%s ..." % message
        self.printSUCCESS(message,printBR=False)
        self.startMessage = message

    def printByResult(self,result,failMessage=None):
        """Print message and result"""
        if self.startMessage:
            self.defaultPrint("\r")
            self.printLine((('greenBr',' * '),
                            ('',self.startMessage),
                           ),
                           (('',''),
                           ), 0, False)
            self.printRight(self.lenString(self.startMessage)+5,4)
            if result:
                self.printLine((('blueBr','['),
                                ('greenBr',' ok '),
                                ('blueBr',']')),[])
            else:
                self.printLine((('blueBr','['),
                                ('redBr',' !! '),
                                ('blueBr',']')),[])
                if failMessage:
                    self.printERROR(failMessage)
            self.startMessage = ""

    def _getCommand(self,commandlist):
        return " ".join(map(lambda x:'"%s"'%x if " " in x else x,commandlist))

    def runChroot(self,chroot,command):
        """Run command in chroot specified by cl_assemble_path"""
        try:
            envdict = {'CHROOT':"on"}
            envdict.update(os.environ)
            commandLine = ["chroot",chroot,
                                "/bin/bash","-c",command]
            chrootCommand = process(*commandLine,stderr=PIPE,
                                stdout=PIPE,envdict=envdict)
        except KeyboardInterrupt:
            chrootCommand.kill()
            raise BuilderError(
                _("An interrupt occurred during the command executing")+
                  ":\n    %s"%self._getCommand(chrootCommand.command))
        return chrootCommand

    def cleanNeedlessKernelData(self, distrPath):
        """Cleaning needless kernel files from /boot,/lib/modules,/usr/src"""
        self.printMessageForTest(_("Cleaning %s directory")% "/boot")
        bootDir = path.join(distrPath, "boot")
        modulesDir = path.join(distrPath, "lib/modules")
        reTrash =re.compile("^(?:config|initramfs|vmlinuz|System.map)")
        rightKernelFiles = \
            [self.clVars.Get('cl_builder_initrd_install'),
             self.clVars.Get('cl_builder_kernel'),
             self.clVars.Get('cl_builder_kernel_config'),
             self.clVars.Get('cl_builder_kernel_systemmap')]
        delFiles = filter(lambda x: not x in rightKernelFiles,
                          filter(reTrash.search,
                                 os.listdir(bootDir)))
        map(lambda x:os.unlink(path.join(bootDir,x)),delFiles)
        self.printByResult(True)
        self.printMessageForTest(_("Cleaning %s directory") % "/lib/modules")
        reKernelVer = re.compile(" version (\S+)\s",re.S)
        ftype = typeFile(magic=0x4).getMType
        resReg = reKernelVer.search(ftype(path.join(bootDir,
                    self.clVars.Get('cl_builder_kernel'))))
        if resReg:
            kernelVersion = resReg.groups()[0]
            map(lambda x:removeDir(path.join(modulesDir,x)),
                filter(lambda x:x != kernelVersion,
                       os.listdir(modulesDir)))
            self.printByResult(True)
            self.printMessageForTest(_("Cleaning %s directory")% "/usr/src")
            removeKernelSources = filter(lambda x:x and x != "/usr/src/linux",
                   self.runChroot(distrPath,"qfile -o /usr/src/*"))
            map(lambda x:removeDir(path.join(distrPath,x)),
                removeKernelSources)
            self.printByResult(True)

    def prepareSourceDistributive(self,distr):
        """Unmount all bind,proc mount points from source distribute"""
        mp = self.clVars.Get('cl_builder_path')
        mps = filter(lambda x:x!=mp,map(lambda x:x[1],childMounts(mp)))
        for target in sorted(mps, reverse=True):
            self.printMessageForTest(_("Unmounting %s")%(target[len(mp):]))
            umountProcess = process("umount",target)
            if umountProcess.failed():
                raise BuilderError(_("Can not umount %s")%target)
            self.printByResult(True)
        distrPath = distr.convertToDirectory().getDirectory()
        self.cleanNeedlessKernelData(distrPath)
        self.applyTemplatesForSquash(distrPath)

    def isoPrepacking(self,directory):
        self.printByResult(True)
        self.applyTemplatesForIso(directory)
        self.printMessageForTest(_("Releasing source data"))
        self.sourceDistr.close()
        self.printByResult(True)
        self.sourceDistr = None
        self.printMessageForTest(_("Creating iso image"))

    def flashPrepacking(self,directory):
        self.printByResult(True)
        self.applyTemplatesForIso(directory)
        self.printMessageForTest(_("Releasing source data"))
        self.sourceDistr.close()
        self.sourceDistr = None

    def getFreeFor(self,pathname):
        """Get free for pathname"""
        dfProcess = process("/bin/df","-h",pathname)
        data = dfProcess.readlines()
        if len(data)>1:
            data = filter(lambda x:x,"\t".join(data[1:]).split())
            if len(data)>3:
                return data[3]
        return "Unknown"

    def getParentMountFor(self,pathname):
        dfProcess = process("/bin/df","-h",pathname)
        data = dfProcess.readlines()
        if len(data)>1:
            data = filter(lambda x:x,"\t".join(data[1:]).split())
            if len(data)>5:
                return data[5]
        return ""

    def printInfo(self):
        self.printSUCCESS(_("Creating image of") + " Calculate Linux")
        self.defaultPrint("%s\n"%_("System information"))
        subname = self.clVars.Get('os_builder_linux_subname')
        subname = (" %s"%subname) if subname else ""
        self.printSUCCESS(_("Building system")+": %s"%
                self.clVars.Get('os_builder_linux_name')+subname)
        self.printSUCCESS(_("System version")+": %s"%
                self.clVars.Get('os_builder_linux_ver'))
        self.printSUCCESS(_("Machine hardware name")+": %s"%
                self.clVars.Get('os_builder_arch_machine'))
        self.printSUCCESS(_("Free disk space for iso building: %s")%
            self.getFreeFor(path.dirname(path.normpath(
                    self.clVars.Get('cl_builder_iso_path')))))
        self.printSUCCESS(_("Free disk space for iso image: %s")%
                self.getFreeFor(path.dirname(path.normpath(
                    self.clVars.Get('cl_builder_image')))))
        self.defaultPrint("%s\n"%_("Perform pre-install checkups"))

        imagefile = self.clVars.Get('cl_builder_image')
        if imagefile:
            self.printSUCCESS(_("Image will be created at: %s")% imagefile)
            if path.exists(imagefile):
                self.printWARNING(_("WARNING") +": " +
                    _("image already exists") + ", "+
                    _("continuation of the operation will overwrite it"))
            if self.getParentMountFor(path.dirname(path.normpath(
                    imagefile))) == "/":
                if self.clVars.Get('os_root_type')=='livecd':
                    self.printWARNING(_("WARNING") +": " +
                        _("image will be created on temporary filesystem"))
                else:
                    self.printWARNING(_("WARNING") +": " +
                     _("image will be created on disk which mounted to root"))

            self.defaultPrint("\n")
        else:
            self.printWARNING("No path for image creating.")

    def printRescratchInfo(self):
        self.printSUCCESS(_("Rebuilding live image of") + " Calculate Linux")
        self.defaultPrint("%s\n"%_("System information"))
        subname = self.clVars.Get('os_builder_linux_subname')
        subname = (" %s"%subname) if subname else ""
        self.printSUCCESS(_("Building system")+": %s"%
                self.clVars.Get('os_builder_linux_name')+subname)
        self.printSUCCESS(_("System version")+": %s"%
                self.clVars.Get('os_builder_linux_ver'))
        self.printSUCCESS(_("Machine hardware name")+": %s"%
                self.clVars.Get('os_builder_arch_machine'))
        self.printSUCCESS(_("Free disk space on flash: %s")%
                self.getFreeFor(path.normpath(
                    self.clVars.Get('cl_builder_iso_path'))))
        self.defaultPrint("%s\n"%_("Perform pre-install checkups"))

        self.printSUCCESS(_("Image will be created in: %s")%
            self.clVars.Get('cl_builder_iso_path'))
        if self.clVars.Get('cl_builder_remove_squash'):
            self.printSUCCESS(
                _("Follow squash images will be removed: %s\n")%
                 ", ".join(self.clVars.Get('cl_builder_remove_squash')))
        else:
            self.defaultPrint("\n")

    def checkVariables(self,rescratch=False):
        """Check values of variables"""
        buildDirectory = self.clVars.Get('cl_builder_iso_path')
        if len(filter(lambda x:"cl-builder" in x,getRunCommands()))>2:
            self.printERROR(
                _("Before proceeding, complete the program cl-builder"))
            return False
        if path.realpath(self.clVars.Get('cl_builder_path')) == "/":
            self.printERROR(_("Source system should not be '/'"))
            return False
        minver = "10.8"
        if getTupleVersion(self.clVars.Get('os_builder_linux_ver')) < \
            getTupleVersion(minver):
            self.printERROR(_("Command is supported for system not less version"
                    " %s") % minver )
            return False
        elif not self.clVars.Get('os_builder_linux_shortname') in \
            varsShare.dictNameSystem.keys():
            self.printERROR(_("Distributive is not Calculate Linux"))
            return False
        elif not self.clVars.Get('cl_builder_kernel'):
            self.printERROR(_("Can not detect kernel"))
            return False
        elif not self.clVars.Get('cl_builder_initrd_install'):
            self.printERROR(_("Can not detect initramfs"))
            return False
        elif not rescratch and path.exists(buildDirectory):
            self.printWARNING(
                _("Directory '%s' for building iso is already exists")%
                    buildDirectory)
            return False
        return True

    def _isoLogic(self,force):
        isoFile = self.clVars.Get('cl_builder_image')
        buildDirectory = self.clVars.Get('cl_builder_iso_path')
        sourceDirectory = self.clVars.Get('cl_builder_path')
        bindDirectory = pathJoin(buildDirectory,
                        self.clVars.Get('cl_builder_squash_path'))
        self.printInfo()
        if not self.checkVariables():
            return False
        self.targetDistr = IsoDistributive(isoFile,bdirectory=buildDirectory)
        self.sourceDistr = DirectoryDistributive(sourceDirectory,
                                            mdirectory=bindDirectory)
        if not force:
            dialogMessage = _("Continue the creating image of the system") \
                            + " (yes/no)"
            dialogRes = dialogYesNo(dialogMessage)
            if dialogRes in (None,False):
                self.assembleIso = False
                self.printERROR(_("Interrupting the image creating"))
                return False
        self.printMessageForTest(_("Prepare data for live image"))
        self.prepareSourceDistributive(self.sourceDistr)
        self.printByResult(True)
        self.targetDistr.prepareIso = self.isoPrepacking
        self.printMessageForTest(_("Creating squash image"))
        self.targetDistr.installFrom(self.sourceDistr)
        self.printByResult(True)
        self.printMessageForTest(_("Creating installed package list"))
        self.printByResult(self.createPackageList(sourceDirectory,isoFile+".list"))
        self.printMessageForTest(_("Creating DIGESTS file"))
        self.printByResult(self.createDigest(isoFile,isoFile+".DIGESTS"))

        #self.targetDistr = PartitionDistributive(rootPartdev,flagRemoveDir=False)
        #dd = DirectoryDistributive(mp,parent=self.targetDistr)
        #self.removeVars('os_assemble_linux_ver', 'os_assemble_linux_shortname',
        #                'cl_assemble_image', 'cl_assemble_step_system',
        #                'cl_assemble_make', 'cl_assemble_step_world',
        #                'cl_assemble_step_newuse', 'cl_assemble_step_update',
        #                'cl_assemble_path')
        return True

    def removeVars(self,*varsname):
        """Remove vars specified by varsname"""
        res = True
        for var in varsname:
            res = res and self.clVars.Delete(var,True,'default','assemble')
        return bool(res)

    def _rescratchLogic(self,force=False):
        self.clVars.Set('cl_builder_iso_path','/mnt/flash',True)
        self.clVars.Get('cl_builder_remove_squash')
        buildDirectory = self.clVars.Get('cl_builder_iso_path')
        sourceDirectory = self.clVars.Get('cl_builder_path')
        bindDirectory = pathJoin(buildDirectory,
                        self.clVars.Get('cl_builder_squash_path'))
        if not (self.clVars.Get('os_root_type') == 'livecd' \
            and os.access(self.clVars.Get('cl_builder_iso_path'),os.W_OK)):
            self.printERROR(
                _("The computer must be load from the flash in Builder mode"))
            return False
        self.printRescratchInfo()
        if not self.checkVariables(rescratch=True):
            return False
        self.targetDistr = IsoDistributive(buildDirectory,
                                           bdirectory=buildDirectory)
        self.sourceDistr = DirectoryDistributive(sourceDirectory,
                                            mdirectory=bindDirectory)
        if not force:
            dialogMessage = _("Continue the rebuilding flash live system") \
                            + " (yes/no)"
            dialogRes = dialogYesNo(dialogMessage)
            if dialogRes in (None,False):
                return True
        self.printMessageForTest(_("Prepare data for live image"))
        self.prepareSourceDistributive(self.sourceDistr)
        self.printByResult(True)
        self.targetDistr.prepareIso = self.flashPrepacking
        self.printMessageForTest(_("Creating squash image"))
        self.targetDistr.installFrom(self.sourceDistr)
        self.printByResult(True)

        self.printMessageForTest(_("Removing old images"))
        oldImages = map(lambda x:pathJoin(buildDirectory,x),
                        self.clVars.Get('cl_builder_remove_squash'))
        try:
            map(lambda x:os.unlink(x),oldImages)
        except (Exception,KeyboardInterrupt),e:
            raise BuilderError(_("Can not remove old files")+":\n%s"%str(e))
        self.printByResult(True)

        self.printMessageForTest(_("Performing syncronization"))
        processSync = process("/bin/sync")
        self.printByResult(processSync.success())
        return True

    def makeIsoImage(self,force=False):
        return self.makeImage(self._isoLogic,force)

    def makeRescratch(self,force=False):
        return self.makeImage(self._rescratchLogic,force)

    def makeImage(self,logicFunc,*argv):
        """Make iso image by variables"""
        self.sourceDistr = None
        self.targetDistr = None
        error = None
        try:
            if not logicFunc(*argv):
                return False
        except (EOFError), e:
            error = e
        except (BuilderError,DistributiveError),e:
            error = e
        except (Exception),e:
            error = ""
            for i in apply(traceback.format_exception, sys.exc_info()):
                error += i
        except KeyboardInterrupt,e:
            self.printByResult(False)
            self.printWARNING("Interrupting the iso building")
            error = _("Iso building manually interrupt")
        if error:
            self.printByResult(False)
        try:
            if self.clTempl:
                self.clTempl.closeFiles()
            if self.sourceDistr:
                self.printMessageForTest(_("Releasing source data"))
                self.sourceDistr.close()
                self.printByResult(True)
            if self.targetDistr:
                self.printMessageForTest(_("Unmount built system volume"))
                self.targetDistr.close()
                self.printByResult(True)
            if self.assembleIso:
                self.restoreMount()
        except (BuilderError,DistributiveError),e:
            error = "%s\n%s" % (str(error),_("Unmounting error"))
        except KeyboardInterrupt,e:
            pass
        if error:
            self.printByResult(False)
        if error:
            for line in filter(lambda x: x,str(error).split('\n')):
                self.printERROR(line)
            self.printERROR(_("System building failed"))
            return False
        self.printSUCCESS(_("System has built successfully"))
        return True

    def restoreMount(self):
        """Mount /proc,/sys, remote, /dev to chroot"""
        mountResources = (("/dev",None,None),
                          ("/dev/shm",None,None),
                          ("/dev/pts",None,None),
                          (None,"-t proc","/proc"),
                          (None,"-t sysfs","/sys"),
                          ("/var/calculate/remote",None,None))
        assemblePath = self.clVars.Get('cl_builder_path')
        for source,opts,target in mountResources:
            opts = opts or "-o bind"
            target = target or source
            self.printMessageForTest(_("Mounting %s")%(source or target))
            target = pathJoin(assemblePath,target)
            if not path.exists(target):
                os.makedirs(target,mode=0755)
            if source == isMount(target):
                continue
            args = ["mount"]+opts.split()+[str(source).lower(),target]
            mountProcess = process(*args)
            if mountProcess.failed():
                raise BuilderError(_("Can not mount %(from)s to %(to)s")%
                        {'from':source,'to':target})
            self.printByResult(True)

    def setAssembleData(self,newprofile):
        """Get assemble data from assemble.env"""
        envData = iniParser(self.envFile)
        if newprofile == None:
            distros = self.clVars.Get('cl_builder_distro')
            if len(distros) == 1:
                newprofile = distros[0]
        newprofile = newprofile or ""
        reProfile = re.compile('(^|/)%s(/|$)'%newprofile.strip('/'),re.S)
        likeProfile = filter(reProfile.search,
                             self.clVars.Get('cl_builder_distro'))
        if len(likeProfile) != 1:
            if newprofile != "list":
                if not newprofile:
                    self.printERROR(_("need specify '--profile'"))
                    return False
                elif not likeProfile:
                    self.printERROR(_("wrong value for '--profile'"))
                    self.printERROR(
                        _("specify 'list' in '--profile' value for print "
                          "all available assembled distro"))
                    return False
                else:
                    self.printERROR(
                        _("specified value of '--profile' is ambiguous. "
                        "Please specify profile more exactly."))
                    self.defaultPrint(_("Select profile from")+":\n")
            else:
                likeProfile = self.clVars.Get('cl_builder_distro')
                self.defaultPrint(_("Available assembled distro")+":\n")
            for profile in likeProfile:
                self.printSUCCESS(profile)
            return False
        newprofile = likeProfile[0]
        self.clVars.Set('os_builder_profile', newprofile, True)
        self.assembleIso = True
        self.clVars.Set('cl_builder_path', _toUNICODE(envData.getVar(newprofile,
                          'cl_assemble_path')).encode('utf-8'),True)
        linuxver = _toUNICODE(envData.getVar(newprofile,
                          'os_assemble_linux_ver')).encode('utf-8')
        if not linuxver:
            curdate = datetime.now()
            linuxver =  "%04d%02d%02d"%(curdate.year,curdate.month,curdate.day)
        self.clVars.Set('os_builder_linux_ver',linuxver ,True)
        return True

    def createPackageList(self,chrootdir,filename):
        """Create package list by chrootdir"""
        pkgdir = path.join(chrootdir,'var/db/pkg')
        if not path.exists(chrootdir):
            return False
        try:
            packageList = sorted(reduce(lambda x,y:x+map(lambda x:path.join(y,x),
                os.listdir(path.join(pkgdir,y))), os.listdir(pkgdir),[]))
            open(filename,'w').writelines(map(lambda x:"%s\n"%x,packageList))
        except (IOError,OSError),e:
            return False
        return True

    def createDigest(self,filename,digestfile):
        """Create file with digest"""
        template = """# %(alg)s HASH\n%(digest)s  %(filename)s\n"""
        try:
            open(digestfile,'w').writelines(map(lambda x:template%{
                'alg':x.upper(),
                'digest': \
                    getattr(hashlib,x)(open(filename,'r').read()).hexdigest(),
                'filename':path.basename(filename)},["md5","sha1"]))
        except (IOError,OSError),e:
            return False
        return True