calculate-utils-2.2-install/pym/cl_distr.py

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

from os.path import exists as pathexists
from os.path import dirname as pathdirname
from os.path import join as pathjoin
from random import choice
import string
import os
import types
from time import sleep
import re
import sys
from cl_utils import runOsCommand,isMount,removeDir,typeFile
from shutil import copyfile
from cl_template import _terms
from subprocess import Popen,PIPE


class DistributiveError(Exception):
    """Error for distributive operations"""
    pass

class DistributiveRepository:
    system = {
            'cds':'server',
            'cld':'desktop',
            'cldg':'desktop',
            'cldx':'desktop',
            'cls':'desktop',
            'css':'server'
            }

    marches = ['i686','x86_64']

    extensiton = ['iso','tar.bz2','tar.gz','tar.7z','tar.lzma']

    reDistName = re.compile("""
                ^(?P<name>%(name)s)
                -(?P<ver>%(ver)s)
                -(?P<march>%(march)s)
                .(?P<ext>%(ext)s)$""" %
                {'name':"|".join(system.keys()),
                 'ver':r"(\d+\.)*\d+",
                 'march':"|".join(marches),
                 'ext':"|".join(extensiton)
                }, re.X)
    def __init__(self,directory):
        self.root = directory

    def _getAvailableDistributives(self,system=None,shortname=None,march=None,
                            version=None):
        """Get all distributives by filter"""
        def distfilter(dist):
            match = self.reDistName.match(dist)
            if not match:
                return False
            d = match.groupdict()
            # check filter conditions
            if system and self.system(d['name']) != system or \
                shortname and d['name'] != shortname or \
                march and d['march'] != march or \
                version and d['ver'] != version:
                return False
            else:
                return True
        return [ dist for dist in os.listdir(self.root) if distfilter(dist) ]

    def _getDistributiveByFile(self,filename):
        """Get Distributive object by filename"""
        # MAGIC_COMPRESS 0x000004 Check inside compressed files 
        tf = typeFile(magic=0x4)
        ftype = tf.getMType(filename)
        if "ISO 9660 CD-ROM" in ftype:
            return IsoDistributive(filename)
        elif "7-zip" in ftype or \
            "POSIX tar archive" in ftype:
            return ArchiveDistributive(filename)
        elif "Squashfs filesystem" in ftype:
            return SquashDistributive(filename)
        else:
            raise DistributiveError("Wrong distributive '%s':\n%s"%
                    (filename,ftype))


    def getLastDistributive(self,system=None,shortname=None,march=None,
                            version=None):
        """Get the actualest distributive"""
        availDistrs = self._getAvailableDistributives(system,shortname,
                                                        march,version)
        # TODO: improve selecting best distributive
        if availDistrs:
            term = _terms()
            bestDistr = availDistrs[0]
            bestDict = self.reDistName.match(bestDistr).groupdict()
            for dist in availDistrs[1:]:
                match = self.reDistName.match(dist)
                d = match.groupdict()
                v1,v2 = term._convertVers(d['ver'],bestDict['ver'])
                if v1 > v2 or \
                    bestDict['ext'] != 'iso' and d['ext'] == 'iso':
                    bestDistr = dist
                    bestDict = d
            return self._getDistributiveByFile(
                            pathjoin(self.root,bestDistr))
        else:
            return None

class Distributive(object):
    """Distributive object. Parent object for all distributive."""
    mountError = "Cann't mount %s:\n%s"
    def __init__(self, parent=None):
        self.childs = []
        # if specified parent
        if parent:
            # save parent type for correct resource release
            self.parent = type(parent)
            # append this object as child to specified parent
            parent.childs.append(self)
        else:
            self.parent = None

    def detach(self):
        """Detach child distributive from parent.
        
        At example: ArchiveDistributive create child DirectoryDistributive by
        unpacking to temporary directory and at close method remove it. If the
        removing directroy do not need, then need call detach in
        DirectoryDistributive object
        
        dist1 = ArchiveDistributive(file="/os.tar.bz2",mdirectory="/tmp/test")
        dist2 = dist1.convertToDirectory()
        dist2.detach()
        dist1.close()
        ...
        """
        self.parent = None

    def close(self):
        """Release all child distributive and release himself.
        
        Need call this method at end work with object for right releasing
        resources.
        
        Example:
        dist1 = PartitionDistributive(partition="/dev/sda2")
        dist2 = dist1.convertToDirectory()
        dist1.close()
        """
        # close all child
        if self.childs:
            for child in self.childs:
                # check detach
                if child.parent:
                    child.close()
            self.childs = []
        # if has parent
        if self.parent:
            self.parent(None).releaseChild(self)
            self.parent = None

    def releaseChild(self,child):
        """Method of release child state of distributive
        
        At example: PartitionDistributive may be DirectoryDistributive by
        mounting it to directory. But at end this directory must be
        unmounted."""
        pass

    def convertToDirectory(self):
        """Default c raise error about impossible convert object"""
        raise DistributiveError("Cann't convert '%s' to '%s'" %
                (self.__class__.__name__,
                 "DirectoryDistributive"))

#    def __del__(self):
#        """Uncomment this method for automaticaly release all distributive
#        instance"""
#        self.close()

    def runOsCommand(self, *argv, **kwarg):
        res,mes = runOsCommand(*argv,**kwarg)
        if type(mes) == types.ListType:
            mes = "\n".join(map(lambda x: x.strip(), mes))
        return res,mes

    def getDirectory(self):
        return self.convertToDirectory().directory

    def _makeDirectory(self,path):
        """Make directory and parent.

        If directory exists then return False else True"""
        try:
            parent = os.path.split(path)[0]
            if not pathexists(parent):
                self._makeDirectory(parent)
            else:
                if pathexists(path):
                    return False
            os.mkdir(path)
            return True
        except (Exception,KeyboardInterrupt),e:
            raise DistributiveError("Cann't create directory '%s':\n%s"%
                        (path,str(e)))

    def _removeDirectory(self,directory):
        """Remove directory and files contained in it"""
        #print "removeDirectory:",directory
        try:
            removeDir(directory)
        except (Exception,KeyboardInterrupt),e:
            raise DistributiveError("Unable remove directory from '%s':\n%s"
                        %(directory,str(e)))

    def _copyfile(self,infile,outfile):
        try:
            copyfile(infile,outfile)
        except (Exception,KeyboardInterrupt),e:
            raise DistributiveError("Cann't copy '%s' to '%s':\n%s"%
                    (infile,outfile,str(e)))

    def copy(self,fromdir,todir):
        """Copy files from 'fromdir' directory to 'todir' directory"""
        res,errmes = self.runOsCommand("rsync -a -x %s/ %s/" % (fromdir,todir))
        if res == 0:
            return True
        else:
            raise DistributiveError("Cann't copy file from '%s' to '%s':\n%s"%
                    (fromdir,todir,errmes))

    def installBootloader(self):
        pass

    def formatPartition(self,format=""):
        pass

    def rndString(self):
        """Get random string with len 8 char"""
        return "".join([choice(string.ascii_letters+string.digits)
                    for i in xrange(0,8)])

    def _mountToDirectory(self,file,directory,mountopts=""):
        """Mount squashfs to directory"""
        if isMount(directory):
            raise DistributiveError("Cann't mount to directory:\n"+
                            "Directory already mounted")
        res,errmes = self.runOsCommand("/bin/mount %s %s %s"%
                                (mountopts,file,directory))
        if res == 0:
            return True
        else:
            try:
                self._removeDirectory(directory)
            except:
                pass
            raise DistributiveError(self.mountError%(file,errmes))

    def _umountDirectory(self,directory):
        """Umount directory"""
        if isMount(directory):
            for wait in [0,0.5,2,5]:
                sleep(wait)
                res,errmes = self.runOsCommand("/bin/umount %s"%directory)
                if res == 0:
                    return True
            raise DistributiveError("Cann't umount %s:\n%s"%
                                    (directory,errmes))
        else:
            return True

    def _getMntDirectory(self,directory):
        """Get directory name, which will use for mounting or unpacking
        
        If queried name is not free then to name append random string
        """
        newDirectoryName = directory
        while pathexists(newDirectoryName):
            newDirectoryName = "%s.%s"%(directory,self.rndString())
        return newDirectoryName

class DirectoryDistributive(Distributive):
    def __init__(self,directory,parent=None):
        Distributive.__init__(self,parent=parent)
        self.directory = directory
        if not parent:
            self._makeDirectory(self.directory)

    def convertToDirectory(self):
        return self

    def installFrom(self, source):
        """Install distributive to directory from source distributive"""
        if isinstance(source,ArchiveDistributive):
            source.unpackTo(self.directory)
        else:
            # get source distributive as directory distributive
            dFrom = source.convertToDirectory()
            # copy distributive from source to this
            self.copy(dFrom.directory,self.directory)

class PartitionDistributive(Distributive):
    reRightPartition = re.compile(r"^/dev/(sd[a-z]\d+|cciss/c\d+d\d+)$")

    def __init__(self,partition,parent=None,mdirectory="/mnt/calculate",
                    check=False,multipartition=None):
        """Initialize partition distributive

        mdirectory - directory for mount
        check - check partition name and raise DistributiveError if partition
                has bad name
        """
        Distributive.__init__(self,parent=parent)
        self.partition = partition
        self.mdirectory = mdirectory
        if check and not partition is None and \
            (not self.reRightPartition.match(partition) or \
            not pathexists(partition)):
            raise DistributiveError("Wrong partition name '%s'"%partition)

    def _mountPartition(self,partition,directory):
        """Mount partition to directory"""
        self._makeDirectory(directory)
        self._mountToDirectory(partition,directory)

    def _umountPartition(self,directory):
        """Umount partition and remove directory"""
        self._umountDirectory(directory)
        self._removeDirectory(directory)

    def releaseChild(self,child):
        """Umount child Directory distributive"""
        # TODO: need multipartition unmounting
        if isinstance(child,DirectoryDistributive):
            self._umountPartition(child.directory)
            child.directory = None

    def convertToDirectory(self):
        """Convert partition to directory by mounting"""
        # TODO: need multipartition mounting
        mdirectory = self.mdirectory
        for child in self.childs:
            if isinstance(child,DirectoryDistributive) and \
                mdirectory in child.directory:
                return child
        mdirectory = self._getMntDirectory(mdirectory)
        self._mountPartition(self.partition,mdirectory)
        return DirectoryDistributive(mdirectory,parent=self)

    def formatPartition(self,format=""):
        """Format partition"""
        # TODO: need multipartition formating
        res,errmes = self.runOsCommand("/sbin/mkfs.reiserfs -f %s" % self.partition)
        if res == 0:
            return True
        else:
            raise DistributiveError("Cann't format partition %s:\n%s"%
                    (self.partition,errmes))

    def installBootloader(self):
        # TODO: need detect HP hdd which has name c[0-9]d[0-9]
        rePartition = re.compile("sd([a-z])(\d+)",re.S)
        res = rePartition.search(self.partition)
        if res:
            hddNum = ord(res.groups()[0])-ord('a')
            partitionNum = int(res.groups()[1])
        else:
            raise DistributiveError("Wrong partition name '%s'"%self.partition)
        pipe = Popen(["/sbin/grub",
               "--device-map=%s/boot/grub/device.map"%self.getDirectory(),
               "--batch"], stdin=subprocess.PIPE,stdout=subprocess.PIPE)
        pipe.stdin.write("root (hd%d,%d)\n"%(hddNum,partitionNum))
        # TODO: change hd0 for bootloader install to other disks
        pipe.stdin.write("setup (hd0)\n")
        pipe.stdin.write("quit\n")
        pipe.stdin.close()
        # TODO: break if command is hang
        if pipe.wait() != 0:
            raise DistributiveError("Cann't install bootloader")

    def installFrom(self, source):
        """Install distributive to partition from source distributive"""
        # get currect partition as directory
        distrTo = self.convertToDirectory()
        # install into directroy distributive from source
        distrTo.installFrom(source)

class ArchiveDistributive(Distributive):
    def __init__(self,file,parent=None,mdirectory="/var/calculate/tmp/stage"):
        Distributive.__init__(self,parent=parent)
        self.file = file
        self.mdirectory = mdirectory

    def _detectArchive(self,file):
        """Detect archive by "/usr/bin/file" command

        Return bzip2,gzip,7z or None
        """
        res,mes = self.runOsCommand("/usr/bin/file %s"%file)
        if "bzip2 compressed data" in mes:
            return "bzip2"
        elif "gzip compressed data" in mes:
            return "gzip"
        elif "7-zip archive data" in mes:
            return "7z"
        else:
            return None

    def _unpackArchive(self,archfile,directory):
        """Unpack archive"""
        # archive is exists
        if not pathexists(archfile):
            raise DistributiveError("File '%s' not found"%archfile)
        # detect type archive
        archiveType = self._detectArchive(archfile)
        # make directory if archive was detected normally
        if archiveType:
            self._makeDirectory(directory)
        # unpack archive
        if archiveType == "7z":
            res,mes = self.runOsCommand("7za x -so %s | tar xf - -C %s/"%
                (archfile,directory))
        elif archiveType == "bzip2":
            res,mes = self.runOsCommand("tar xjf %s -C %s/"%
                (archfile,directory))
        elif archiveType == "gzip":
            res,mes = self.runOsCommand("tar xf %s -C %s/"%
                (archfile,directory))
        else:
            raise DistributiveError("Unknown archive type by '%s'"%file)
        if res != 0:
            raise DistributiveError("Error during unpacking\n%s"%mes)

    def unpackTo(self,directory):
        """Unpack currect archive to directory"""
        self._unpackArchive(self.file,directory)

    def convertToDirectory(self):
        """Get archive as directory (unpack to directory)"""
        # check may be the archive already unpacked
        mdirectory = self.mdirectory
        for child in self.childs:
            if isinstance(child,DirectoryDistributive) and \
                mdirectory in child.directory:
                return child
        # get temporary directory for unpacking
        mdirectory = self._getMntDirectory(mdirectory)
        dirdist = DirectoryDistributive(mdirectory,parent=self)
        self._unpackArchive(self.file,mdirectory)
        return dirdist

    def releaseChild(self,child):
        """Remove child Directory distributive"""
        if isinstance(child,DirectoryDistributive):
            self._removeDirectory(child.directory)
            child.directory = None

    def packToArchive(self,directory,file):
        res,errmes = self.runOsCommand("tar cf %s -C %s ."%(file,directory))
        if res != 0:
            raise DistributiveError("Cann't create archive '%s':\n%s"%
                                    (file,errmes))

    def installFrom(self, source):
        """Install distributive to partition from source distributive"""
        # get source distributive as directory distributive
        dFrom = source.convertToDirectory()
        # install into directroy distributive from source
        self.packToArchive(dFrom.directory, self.file)

class SquashDistributive(Distributive):
    def __init__(self,file,parent=None,mdirectory="/mnt/livecd"):
        Distributive.__init__(self,parent=parent)
        self.file = file
        self.mdirectory = mdirectory

    def _mountSquash(self,file,directory):
        """Mount squashfs to directory"""
        self._makeDirectory(directory)
        self._mountToDirectory(file,directory,mountopts="-o loop -t squashfs")

    def _umountSquash(self,directory):
        self._umountDirectory(directory)
        self._removeDirectory(directory)

    def convertToDirectory(self):
        mdirectory = self.mdirectory
        for child in self.childs:
            if isinstance(child,DirectoryDistributive) and \
                mdirectory in child.directory:
                return child
        mdirectory = self._getMntDirectory(mdirectory)
        self._mountSquash(self.file,mdirectory)
        return DirectoryDistributive(mdirectory,parent=self)

    def releaseChild(self,child):
        """Umount child Directory distributive"""
        if isinstance(child,DirectoryDistributive):
            self._umountSquash(child.directory)
            child.directory = None

    def packToSquash(self,directory,file):
        res,errmes = self.runOsCommand("/usr/bin/mksquashfs %s/ %s"%
                                (directory,file))
        if res != 0:
            raise DistributiveError("Cann't create squashfs '%s':\n%s"%
                                (file,errmes))

    def installFrom(self, source):
        """Install distributive to partition from source distributive"""
        # get source distributive as directory distributive
        dFrom = source.convertToDirectory()
        # install into directroy distributive from source
        self.packToSquash(dFrom.directory, self.file)

class IsoDistributive(Distributive):
    reLive = re.compile(r"^live[^.]*(\.(\d+))?\.squashfs$",re.S)

    def __init__(self,file,parent=None,mdirectory="/mnt/cdrom",
            bdirectory="/var/calculate/tmp/iso"):
        Distributive.__init__(self,parent=parent)
        self.file = file
        self.mdirectory = mdirectory
        self.bdirectory = bdirectory

    def _mountIso(self,file,directory):
        self._makeDirectory(directory)
        self._mountToDirectory(file,directory,mountopts="-o loop")

    def _umountIso(self,directory):
        self._umountDirectory(directory)
        self._removeDirectory(directory)

    def _getLastLive(self,directory):
        curfile = None
        curnum = -1
        for file in os.listdir(directory):
            res = self.reLive.search(file)
            if res:
                if res.groups()[1]:
                    num = int(res.groups()[1])
                else:
                    num = 0
                if num > curnum:
                    curfile = file
                    curnum = num
        return curfile

    def convertToSquash(self):
        mdirectory = self.mdirectory
        for child in self.childs:
            if isinstance(child,SquashDistributive) and \
                mdirectory in child.file:
                return child
        mdirectory = self._getMntDirectory(mdirectory)
        self._mountIso(self.file,mdirectory)
        fileLive = self._getLastLive(mdirectory)
        if not fileLive:
            raise DistributiveError("Iso %s doesn't contain live image" %
                    self.file)
        return SquashDistributive(pathjoin(mdirectory,fileLive),
                                    parent=self)

    def releaseChild(self,child):
        """Umount child Directory distributive"""
        if isinstance(child,SquashDistributive):
            self._umountIso(pathdirname(child.file))
            child.directory = None

    def convertToDirectory(self):
        return self.convertToSquash().convertToDirectory()

    def prepareIso(self,directory):
        print("apply iso templates to %s/"%directory)
        self._makeDirectory(pathjoin(directory,"isolinux"))
        self._copyfile("/usr/share/syslinux/isolinux.bin",
            pathjoin(directory,"isolinux/isolinux.bin"))

    def packToIso(self,directory,file):
        # remove previous version of iso
        try:
            if pathexists(file):
                os.unlink(file)
        except (Exception,KeyboardInterrupt),e:
            raise DistributiveError("Cann't remove %s:\n%s"%(file,str(e)))

        res,errmes = self.runOsCommand(
            "%(progname)s %(params)s -o %(target)s %(source)s/"%
            {'progname':'/usr/bin/mkisofs',
             'params':" ".join(["-b isolinux/isolinux.bin","-no-emul-boot",
                    "-boot-load-size 4","-boot-info-table","-iso-level 4",
                    "-hide boot.catalog"]),
             'target':file,
             'source':directory})
        if res == 0:
            return True
        else:
            raise DistributiveError("Cann't create iso image %s:\n%s"%
                        (file,errmes))

    def installFrom(self, source):
        """Install distributive to partition from source distributive"""
        # make temporary directory for creating iso image
        isoDirectory = self.bdirectory
        isoDirectory = self._getMntDirectory(isoDirectory)
        self._makeDirectory(isoDirectory)

        try:
            # getting squash from source
            if isinstance(source,SquashDistributive):
                self._copyfile(source.file,
                    pathjoin(isoDirectory,"livecd.squashfs"))
            else:
                distDirectory = source.convertToDirectory()
                squashDistr = SquashDistributive(
                               pathjoin(isoDirectory,"livecd.squashfs"))
                squashDistr.installFrom(distDirectory)

            # prepare iso
            self.prepareIso(directory)
            # pack iso
            self.packToIso(isoDirectory, self.file)
            # remove temporary directory after creating iso image
            self._removeDirectory(isoDirectory)
        except DistributiveError,e:
            self._removeDirectory(isoDirectory)
            raise e
        except KeyboardInterrupt,e:
            self._removeDirectory(isoDirectory)
            raise DistributiveError("Keyboard interrupt")

class ScratchDistributive(Distributive):
    def __init__(self,directory,parent=None,mdirectory="/mnt/install",
                    check=False):
        Distributive.__init__(self,parent=parent)
        self.directory = directory
        self.mdirectory = mdirectory
        if check and not (pathexists(pathjoin(directory,"workspace")) and \
                          pathexists(pathjoin(directory,"delta")) and \
                          pathexists(pathjoin(directory,"calculate"))):
            raise DistributiveError(
                "Wrong scratch distributive in '%s'"%directory)

    def _mountLayers(self,source,target):
        """Mount squashfs to directory"""
        self._makeDirectory(target)
        self._mountToDirectory("none",target,
            mountopts="-t aufs "+\
            "-o udba=relval,br:%(work)s=rw:%(delta)s=ro+wh:%(static)s=ro" %\
                    {"work":pathjoin(source,"workspace"),
                     "delta":pathjoin(source,"delta"),
                     "static":pathjoin(source,"calculate")})

    def _umountLayers(self,directory):
        self._umountDirectory(directory)
        self._removeDirectory(directory)

    def convertToDirectory(self):
        mdirectory = self.mdirectory
        for child in self.childs:
            if isinstance(child,DirectoryDistributive) and \
                kdirectory in child.directory:
                return child
        mdirectory = self._getMntDirectory(mdirectory)
        self._mountLiveImage(self.file,pathjoin(directory,calculate))
        self._mountLayers(self.file,mdirectory)
        return DirectoryDistributive(mdirectory,parent=self)

    def releaseChild(self,child):
        """Umount child Directory distributive"""
        if isinstance(child,DirectoryDistributive):
            self._umountLayers(child.directory)
            #self._umountLiveImage(
            child.directory = None

    def installFrom(self, source):
        """Install distributive to partition from source distributive"""
        # get source distributive as directory distributive
        dFrom = source.convertToDirectory()
        # install into directroy distributive from source
        self.packToSquash(dFrom.directory, self.file)

class ScratchPartitionDistributive(PartitionDistributive):
    reRightPartition = re.compile(r"^/dev/(sd[a-z]\d+|cciss/c\d+d\d+)$")

    def __init__(self,partition,parent=None,mdirectory="/mnt/scratch",
                check=False,multipartition=None):
        """Initialize partition distributive

        mdirectory - directory for mount
        check - check partition name and raise DistributiveError if partition
                has bad name
        """
        Distributive.__init__(self,parent=parent)
        self.partition = partition
        self.mdirectory = mdirectory
        if check and not partition is None and \
            (not self.reRightPartition.match(partition) or \
            not pathexists(partition)):
            raise DistributiveError("Wrong partition name '%s'"%partition)

    def _mountPartition(self,partition,directory):
        """Mount partition to directory"""
        self._makeDirectory(directory)
        self._mountToDirectory(partition,directory)

    def _umountPartition(self,directory):
        """Umount partition and remove directory"""
        self._umountDirectory(directory)
        self._removeDirectory(directory)

    def releaseChild(self,child):
        """Umount child Directory distributive"""
        if isinstance(child,ScratchDistributive):
            self._umountPartition(child.directory)
            child.directory = None

        return None

    def formatPartition(self,format="reiserfs"):
        res,errmes = self.runOsCommand("/sbin/mkfs.reiserfs -f %s" % self.partition)
        if res == 0:
            return True
        else:
            raise DistributiveError("Cann't format partition %s:\n%s"%
                    (self.partition,errmes))

    def convertToDirectory(self):
        return self.convertToScratch().convertToDirectory()

    def convertToScratch(self):
        mdirectory = self.mdirectory
        for child in self.childs:
            if isinstance(child,DirectoryDistributive) and \
                mdirectory in child.directory:
                return child
        mdirectory = self._getMntDirectory(mdirectory)
        self._mountPartition(self.partition,mdirectory)
        return ScratchDistributive(mdirectory,parent=self)

    def prepareScratch(self,directory):
        for scrDirectory in ["calculate","delta","workspace"]:
            self._makeDirectory(pathjoin(directory,scrDirectory))

    def installFrom(self, source):
        """Install distributive to partition from source distributive"""
        # get currect partition as directory
        distrTo = self.mountToDirectory()
        scratchDirectory = distrTo().directory

        try:
            # getting squash from source
            if isinstance(source,SquashDistributive):
                self._copyfile(source.file,
                    pathjoin(scratchDirectory,"livecd.squashfs"))
            else:
                distDirectory = source.convertToDirectory()
                squashDistr = SquashDistributive(
                               pathjoin(scratchDirectory,"livecd.squashfs"))
                squashDistr.installFrom(distDirectory)

            # prepare scratch
            self.prepareScratch(directory)
        except DistributiveError,e:
            self._removeDirectory(scratchDirectory)
            raise e
        except KeyboardInterrupt,e:
            self._removeDirectory(scratchDirectory)
            raise DistributiveError("Keyboard interrupt")