#coding:utf-8
import os
import re
import sys
import grp
import stat
import time
import syslog
import tempfile

syslog.openlog('SUNHPC')
def msg(message):
    syslog.syslog('DOP: %s' % message)

class Partition:

    raidinfo    = ''
    mountpoints = []
    saved_fstab = []

    def __init__(self):
        if os.path.exists('/mnt/runtime/usr/sbin/parted'):
            self.parted = '/mnt/runtime/usr/sbin/parted'
        else:
            self.parted = '/sbin/parted'

        if os.path.exists('/mnt/runtime/usr/sbin/e2label'):
            #self.e2label = '/mnt/runtime/usr/sbin/e2label'
            self.xfs_admin = '/mnt/runtime/usr/sbin/xfs_admin'
        else:
            self.xfs_admin = '/sbin/xfs_admin'

        if os.path.exists('/mnt/runtime/usr/sbin/mdadm'):
            self.mdadm = '/mnt/runtime/usr/sbin/mdadm'
        else:
            self.mdadm = '/sbin/mdadm'

        if not os.path.exists('/tmp/discovered.disks'):
            self.discoveredDisks()

    def discoveredDisks(self):
        raids = []
        disks = []
        if not disks:
            disks = Device().get_disks

        if raids:
            mdstat = open('/proc/mdstat')
            for line in mdstat:
                if line.startswith('md'):
                    devicesName = line.split()[4:]
                    for devName in devicesName:
                        devName = devName.strip('0123456789[]')
                        if devName not in disks:
                            disks.append(devName)
            mdstat.close()

        # Print disks
        discoveredDisk = open('/tmp/discovered.disks', 'w')
        discoveredDisk.write("disks: " + ' '.join(disks) + '\n')
        discoveredDisk.write("raids: " + ' '.join(raids) + '\n')
        discoveredDisk.close()

    def getDisks(self):
        disks = []
        fe = open('/tmp/discovered.disks', 'r')
        for line in fe.readlines():
            l = line.split()
            if len(l) > 0 and l[0] == 'disks:':
                for d in l[1:]:
                    #
                    # only include disks that have their
                    # 'media present' -- this is one way
                    # to filter out Dell Virtual Floppy
                    # devices
                    #
                    disks.append(d)
        fe.close()
        return disks

    def getRaids(self):
        raids = []
        file = open('/tmp/discovered.disks', 'r')
        for line in file.readlines():
            l = line.split()
            if len(l) > 0 and l[0] == 'raids:':
                raids = l[1:]
        file.close()
        try:
            file = open('/dev/md/md-device-map', 'r')
            for line in file.readlines():
                l = line.split()
                if len(l) > 0 and not (l[0] in raids):
                    raids.append(l[0])
            file.close()
        except:
            #no raid
            pass    
        return raids

    def gptDrive(self, devname):
        #
        # if this is a drive with a GPT format, then return '1'
        #
        retval = 0

        cmd = '%s /dev/%s print -s 2> /dev/null' % \
            (self.parted, devname)

        label = 'Partition Table:'
        for line in os.popen(cmd).readlines():
            if len(line) > len(label) and \
                line[0:len(label)] == label:

                l = line.split()
                if len(l) > 2 and l[2] == 'gpt':
                    retval = 1
                    break

        return retval


    def getDevice(self, strs):
        device = ''

        a = strs.split('/dev/')
        if len(a) > 1:
            device = a[1]

        return device.strip()


    def getSectorStart(self, str):
        sectorstart = ''

        a = str.split('=')
        if len(a) > 1 and a[0].strip() == 'start':
            sectorstart = a[1]
        else:
            sectorstart = a[0]

        return sectorstart.strip()


    def getPartitionSize(self, str):
        partitionsize = ''

        a = str.split('=')
        if len(a) > 1 and a[0].strip() == 'size':
            partitionsize = a[1]
        else:
            partitionsize = a[0]

        return partitionsize.strip()


    def getPartId(self, str):
        partid = ''

        a = str.split('=')
        if len(a) > 1 and a[0].strip() == 'Id':
            partid = a[1]
        else:
            partid = a[0]
        
        return partid.strip()


    def getFsType(self, mntpoint):
        return self.findFsTypeInFstab(mntpoint)


    def getBootFlags(self, str):
        return str.strip()
        

    def getMountPoint(self, devicename):
        mntpoint = ''

        cmd = 'blkid -o export /dev/%s | grep UUID ' % devicename
        cmd += ' 2> /dev/null'
        uuid = os.popen(cmd).readlines()
        if len(uuid) > 0:
            mntpoint = self.findMntInFstab(uuid[0][:-1])
        
        if mntpoint == '':
            mntpoint = self.findMntInFstab('/dev/' + devicename)

        if mntpoint == '':
            #
            # see if the device is part of a raidset
            #
            mntpoint = self.getRaidName(devicename)

        if mntpoint == '':
            cmd = '%s /dev/%s 2> /dev/null' % \
                (self.xfs_admin, devicename)
            label = os.popen(cmd).readlines()

            label = ''.join(label)
            id = 'LABEL=%s' % (label[:-1])

            mntpoint = self.findMntInFstab(id)

        return mntpoint


    def getRaidName(self, partition_device):
        raidname = ''

        for info in self.raidinfo:
            if len(info) > 3:
                (device, partitions, raidlevel,
                    num_partitions) = info

                if partition_device in partitions:
                    raidname = 'raid.%s' % partition_device
                    break

        return raidname


    def findMntInFstab(self, identifier):
        for line in self.saved_fstab:
            l = line.split()
            if len(l) > 0:
                if l[0] == identifier:
                    return l[1]

        return ''

    def findFsTypeInFstab(self, mntpoint):
        for line in self.saved_fstab:
            l = line.split()
            if len(l) > 2:
                if l[1] == mntpoint:
                    return l[2]

        return ''


    def formatPartedNodePartInfo(self, devname, info):
        #
        # this function parses partition info from 'parted'
        #
        partinfo = []
        isDisk = 0
        
        for line in info:
            l = line[:-1].split()

            if len(l) > 2 and re.match('[0-9]+', l[0]):
                if devname[0:2] == 'md':
                    device = devname
                elif len(devname) > 4 and \
                        devname[0:5] == 'cciss':
                    #
                    # special case for HP smart array
                    # controllers
                    #
                    device = devname + 'p' + l[0]
                else:
                    device = devname + l[0]
                isDisk = 1
            else:
                if len(l) > 1 and l[0] == 'Disk':
                    isDisk = 1
                continue

            sectorstart = l[1]
            partitionsize = l[3]
            partid = ''
            
            if devname[0:2] == 'md' and len(l) > 4:
                #
                # special case for software RAID. there is
                # no 'Type' or 'Flags' fields, so the
                # 'File system' field is 5th field
                #
                fstype = l[4]
                bootflags = ''
            else:
                bfs = None  
                if len(l) > 5 and not self.gptDrive(devname):
                    #
                    # there is a case for RAID 0 that the 
                    # second partition of the drive does not
                    # get the file system label
                    # (e.g., ext4), so the 'boot flags'
                    # get misidentified as a
                    # file system type
                    #
                    if 'raid' in l[5] or 'boot' in l[5]:
                        bfs = l[5:]
                        fstype = ''
                    
                    else:
                        fstype = l[5]
                else:
                    fstype = ''
            
                    if len(l) > 4 and self.gptDrive(devname):
                        # gpt partition there is not 
                        # Type column so fstype is in 4 
                        fstype = l[4]
                        if len(l) > 5:
                            bfs = [l[5]]
                    

                if not bfs and len(l) > 6:
                    bfs = l[6:]

                if bfs:
                    bf = []                                         
                    for b in bfs:
                        bf.append(b.rstrip(','))                
                    bootflags = ' '.join(bf)  
                else:
                    bootflags = ''

            if 'linux-swap' in fstype:
                mntpoint = 'swap'
            else:
                mntpoint = self.getMountPoint(device)

            # print 'formatPartedNodePartInfo:l: ', l

            partinfo.append('%s,%s,%s,%s,%s,%s,%s,%s\n' %
                    (device, sectorstart, partitionsize,
                        partid, fstype, bootflags, '',
                                mntpoint))

        # print 'formatPartedNodePartInfo:partinfo: ', partinfo

        if partinfo == [] and isDisk:
            #
            # this disk has no partitions, create a
            # dummy null entry for it
            #
            partinfo = [ '%s,,,,,,,\n' % (devname) ]

        return partinfo


    def parsePartInfo(self, info):
        n = info.split(',')

        if len(n) != 8:
            return ('', '', '', '', '', '', '', '')

        device          = n[0].strip()
        sectorstart     = n[1].strip()
        partitionsize   = n[2].strip()
        partid          = n[3].strip()
        fstype          = n[4].strip()
        bootflags       = n[5].strip()
        partflags       = n[6].strip()
        mntpoint        = n[7].strip()

        return (device, sectorstart, partitionsize, partid, 
            fstype, bootflags, partflags, mntpoint)


    def getDiskInfo(self, disk):
        syslog.syslog('getDiskInfo: disk:%s' % (disk))

        cmd = '%s /dev/%s ' % (self.parted, disk)
        cmd += 'print -s 2> /dev/null'
        diskinfo = os.popen(cmd).readlines()
        
        syslog.syslog('getNodePartInfo: diskinfo:%s' % (diskinfo))

        return diskinfo


    def getRaidLevel(self, device):
        level = None

        cmd = '%s --query --detail ' % (self.mdadm)
        cmd += '/dev/%s' % (device)
        for line in os.popen(cmd).readlines():
            l = line.split()
            if len(l) > 3 and l[0] == 'Raid' and l[1] == 'Level':
                if l[3][0:4] == 'raid':
                    level = l[3][4:]
                    break
        
        return level


    def getRaidParts(self, device):
        parts = []

        foundparts = 0
        cmd = '%s --query --detail ' % (self.mdadm)
        cmd += '/dev/%s' % (device)
        for line in os.popen(cmd).readlines():
            l = line.split()
            if len(l) > 4 and l[3] == 'RaidDevice':
                foundparts = 1
                continue

            if foundparts == 0:
                continue

            if len(l) == 0:
                continue
            
            part = l[-1].split('/')
            parts.append('raid.%s' % part[-1])

        return ' '.join(parts)

    
    def getNodePartInfo(self, disks):
        arch = os.uname()[4]

        partinfo = []
        nodedisks = {}

        for line in self.getFstab(disks):
            self.saved_fstab.append(line)

        for devname in disks:
            diskinfo = self.getDiskInfo(devname)
            partinfo += self.formatPartedNodePartInfo(devname,
                diskinfo)

        syslog.syslog('getNodePartInfo: partinfo:%s' % (partinfo))

        for node in partinfo:
            n = self.parsePartInfo(node)

            (nodedevice, nodesectorstart, nodepartitionsize,
                nodepartid, nodefstype, nodebootflags,
                nodepartflags, nodemntpoint) = n

            if (len(nodedevice) > 2) and (nodedevice[0:2] == 'md'):
                nodepartflags = '--level=%s' % \
                    self.getRaidLevel(nodedevice)

                nodebootflags = self.getRaidParts(nodedevice)

                n = (nodedevice, nodesectorstart,
                    nodepartitionsize,
                    nodepartid, nodefstype,
                    nodebootflags,
                    nodepartflags, nodemntpoint)

            elif nodebootflags != '':
                if 'raid' in nodebootflags.split():
                    nodemntpoint = 'raid.%s' % (nodedevice)

                    n = (nodedevice, nodesectorstart,
                        nodepartitionsize,
                        nodepartid, nodefstype,
                        nodebootflags,
                        nodepartflags, nodemntpoint)
                
            if nodedevice != '':
                key = ''
                for disk in disks:
                    if len(disk) <= len(nodedevice) and \
                        disk == nodedevice[0:len(disk)]:

                        key = disk
                        break

                if key != '':
                    if key not in nodedisks:
                        nodedisks[key] = [n]
                    else:
                        nodedisks[key].append(n)

        syslog.syslog('getNodePartInfo:nodedisks:%s' % (nodedisks))

        return nodedisks


    def listDiskPartitions(self, disk):
        list = []
        inHeader = 1
        
        if disk[0:2] == 'md':
            return [ (disk, 'dummy') ]

        for part in self.getDiskInfo(disk):
            l = part.split()

            #
            # skip the 'parted' header
            #
            if len(l) > 1 and l[0] == 'Number':
                inHeader = 0
                continue

            if inHeader:
                continue

            partnumber = 0

            #
            # look for a part number
            #
            if len(l) > 2 and re.match('[0-9]+', l[0]):
                partnumber = int(l[0])

            if partnumber > 0:
                if len(disk) > 4 and disk[0:5] == 'cciss':
                    #
                    # special case for HP smart array
                    # controllers
                    #
                    disk = disk + 'p'

                if len(l) > 5:
                    fstype = l[5]
                else:
                    fstype = ''
                    if len(l) > 4 and self.gptDrive( disk ):
                        # this is a gpt partition
                        fstype = l[4]
                        

                list.append(('%s%d' % (disk, partnumber),
                    fstype))

        return list


    def defaultDataDisk(self, disk):
        basename = '/state/partition'
        parts = []

        i = 1
        while 1:
            nextname = '%s%d' % (basename, i)
            if nextname not in self.mountpoints:
                break
            i = i + 1

        p = 'part '
        p += '%s --size=1 ' % (nextname)
        p += '--fstype=ext4 --grow --ondisk=%s ' % (disk)
        self.mountpoints.append(nextname)
        parts.append(p)

        return parts


    def SunhpcGetPartsize(self, mountpoint):
        size = 0

        if mountpoint == 'root':
            size = 60000
        elif mountpoint == 'var':
            size = 10000
        elif mountpoint == 'swap':
            size = 1000

        return size


    def defaultRootDisk(self, disk):
        arch = os.uname()[4]
        parts = []

        if arch == 'ia64':
            p = 'part /boot/efi --size=1000 --fstype=vfat '
            p += '--ondisk=%s\n' % (disk)

        p = 'part '
        p += '/ --size=%d ' % (self.SunhpcGetPartsize('root'))
        p += '--fstype=ext4 --ondisk=%s ' % (disk)
        self.mountpoints.append('/')
        parts.append(p)

        p = 'part '
        p += '/var --size=%d ' % (self.SunhpcGetPartsize('var'))
        p += '--fstype=ext4 --ondisk=%s ' % (disk)
        self.mountpoints.append('/var')
        parts.append(p)

        p = 'part '
        p += 'swap --size=%d ' % (self.SunhpcGetPartsize('swap'))
        p += '--fstype=swap --ondisk=%s ' % (disk)
        self.mountpoints.append('swap')
        parts.append(p)

        parts += self.defaultDataDisk(disk)
        return parts

    def getFstab(self, disks):
        if os.path.exists('/upgrade/etc/fstab'):
            file = open('/upgrade/etc/fstab')
            lines = file.readlines()
            file.close()
            return lines

        #
        # if we are here, let's go look at all the disks for /etc/fstab
        #
        mountpoint = tempfile.mktemp()
        os.makedirs(mountpoint)
        fstab = mountpoint + '/etc/fstab'

        lines = []
        for disk in disks:
            for (partition, fstype) in \
                    self.listDiskPartitions(disk):

                if not fstype or 'linux-swap' in fstype:
                    continue

                os.system('mount /dev/%s %s' \
                    % (partition, mountpoint) + \
                    ' > /dev/null 2>&1')

                if os.path.exists(fstab):
                    file = open(fstab)
                    lines = file.readlines()
                    file.close()

                os.system('umount %s 2> /dev/null' %
                    (mountpoint))

                if len(lines) > 0:
                    break

            if len(lines) > 0:
                break

        try:
            os.removedirs(mountpoint)
        except:
            pass

        return lines

    def isSunhpcDisk(self, partinfo, touchit=0):
        retval = 0
        mountpoint = tempfile.mktemp()
        os.makedirs(mountpoint)
        for part in partinfo:
            (dev,start,size,id,fstype,bootflags,partflags,mnt) = part

            if not fstype or fstype == 'linux-swap':
                continue

            devname = '/dev/%s' % (dev)
            os.system('mount %s %s' % (devname, mountpoint))

            try:
                filename = mountpoint + '/.sunhpc-release'
                if touchit == 1:
                    os.system('touch %s' % filename)

                if os.path.exists(filename):
                    retval = 1
            except:
                pass

            os.system('umount %s' % (mountpoint) + ' > /dev/null 2>&1')
            if retval == 1:
                break
        try:
            os.removedirs(mountpoint)
        except:
            pass
        return retval

    def addPartitions(self, nodepartinfo, format):
        arch  = os.uname()[4]
        parts = []
        #
        # for each partition on a drive, build a partition
        # specification for anaconda
        #
        for node in nodepartinfo:
            if len(node) == 1: continue

            (nodedevice, nodesectorstart, nodepartitionsize,
                nodepartid, nodefstype, nodebootflags,
                    nodepartflags, nodemntpoint) = node

            if arch == 'ia64':
                if nodefstype == 'fat32':
                    nodefstype = 'vfat'
                elif nodefstype == 'linux-swap':
                    nodefstype = 'swap'

            if nodemntpoint == '':
                continue
            #
            # only add raid partitions if they have a mountpoint
            # defined by their respective 'md' device.
            #
            # anaconda will crash if there is not a valid
            # mountpoint for the md device
            #
            if nodepartid == 'fd':
                if not self.getRaidMountPoint(nodedevice):
                    continue
            args = [ nodemntpoint ]
            if len(nodemntpoint) > 3 and \
                        nodemntpoint[0:4] == 'raid':
                #
                # never format a software raid partition and
                # always set its size to 1
                #
                args.append('--noformat')
                args += [ '--size', '1']
            elif (nodemntpoint != '/' and nodemntpoint != '/var') \
                    and not format:
                args.append('--noformat')
            else:
                if nodefstype == '':
                    args += [ '--fstype', self.fstype ]
                else:
                    args += [ '--fstype', nodefstype ]

            israid = 0
            if len(nodedevice) > 2 and nodedevice[0:2] == 'md':
                israid = 1
                args += [ "--device=%s" % (nodedevice) ]
                args += [ "--useexisting" ]

                if nodepartflags != '':
                    args += [ nodepartflags ]
            else:
                args += [ '--onpart', nodedevice ]

            if israid:
                parts.append('raid %s' % (args.join()))
            else:
                parts.append('part %s' % (args.join()))
            
            self.mountpoints.append(nodemntpoint)
        return parts

    def compareDiskInfo(self, dbpartinfo, nodepartinfo):
        if len(dbpartinfo) != len(nodepartinfo):
            return 0

        for db in dbpartinfo:
            if len(db) == 1:
                continue

            (dbdevice, dbsectorstart, dbpartsize, dbpartid,
                dbfstype, dbbootflags, dbpartflags,
                dbmntpoint) = db

            found = 0
            for node in nodepartinfo:
                if len(node) == 1:
                    continue

                (nodedevice, nodesectorstart, nodepartsize,
                    nodepartid, nodefstype, nodebootflags,
                    nodepartflags, nodemntpoint) = node

                # print 'compareDiskInfo:node: ', node
                # print 'compareDiskInfo:db: ', db
                if dbsectorstart == nodesectorstart and \
                    dbpartsize == nodepartsize and \
                    dbpartid == nodepartid and \
                    dbfstype == nodefstype and \
                    dbbootflags == nodebootflags and \
                    dbpartflags == nodepartflags and \
                    dbmntpoint == nodemntpoint:
                    found = 1
                    break
            if not found: return 0
        return 1

class Device:
    def __init__(self, subsys='/sys/class/block'):
        self.subsys = subsys
        self.block  = []
        self.disks  = {}
        self.parts  = {}
        self.build()

    def build(self):
        if not os.path.isdir(self.subsys):
            return self.devices

        for sub in os.listdir(self.subsys):
            self.block.append(Block(sub, self.subsys))

        for disk in self.block:

            if disk.device_is_partition:
                self.parts[disk.info.get('DEVNAME')] = disk

            if not disk.device_is_disk:
                continue

            if disk.device_is_loop:
                continue

            if disk.device_is_dm:
                continue

            self.disks[disk.info.get('DEVNAME')] = disk

    @property
    def get_parts(self):
        return list(self.parts.keys())
    @property
    def get_disks(self):
        return list(self.disks.keys())

    @property
    def get_disk_object(self, dev=None):
        if dev in self.disks:
            return self.disks.get(dev)
    @property
    def get_part_object(self, dev=None):
        if dev in self.parts:
            return self.parts.get(dev)

class Block:
    def __init__(self, devname, block):
        self.devname = devname
        self.block   = block
        self.devinfo = {}
        self.build()

    def build(self):
        self.device_info()

    @property
    def info(self):
        return self.devinfo

    @property
    def device_path(self):
        return os.path.realpath(os.path.join(self.block, self.devname))

    def device_info(self):
        """read /sys/class/block/sda/uevent info"""
        # DEVNAME=sda
        # DEVTYPE=disk
        uevent = os.path.join(self.device_path, 'uevent')
        if not os.path.isfile(uevent):
            return self.devinfo

        with open(uevent) as f:
            for l in f.readlines():
                sp = l.split('=', 1)
                if not len(sp): continue
                self.devinfo[sp[0]] = sp[1].strip()

    @property
    def device_is_cdrom(self):
        try:
            dev = os.path.join(os.sep, 'dev', self.info.get('DEVNAME'))
            gid = os.stat(dev)[stat.ST_GID]
            t = grp.getgrgid(gid).gr_name
        except:
            t = None

        if t and t == 'cdrom':
            return True

        if os.path.exists('/opt/sunhpc/bin/lsscsi'):
            lsscsi = '/opt/sunhpc/bin/lsscsi'
        else:
            lsscsi = '/sbin/lsscsi'

        if not os.path.exists(lsscsi):
            return False

        line = os.popen('%s |grep "%s"' % (lsscsi, self.info.get('DEVNAME'))).readline().split()
        if len(line) > 4:
            if line[1] == 'cd/dvd' or line[3] == 'DVD-ROM':
                return True
        return False

    @property
    def device_is_disk(self):
        if self.device_is_cdrom:
            return False
        has_range = os.path.exists(os.path.join(self.device_path, 'range'))
        return self.info.get("DEVTYPE") == "disk" or has_range

    @property
    def device_is_partition(self):
        has_start = os.path.exists(os.path.join(self.device_path, 'start'))
        return self.info.get("DEVTYPE") == "partition" or has_start

    @property
    def device_is_loop(self):
        return self.info.get('DEVNAME').startswith('loop')

    @property
    def device_is_dm(self):
        return self.info.get('DEVNAME').startswith('dm')

    def __repr__(self):
        return 'Device("%s")' % self.device_path
