#!/usr/bin/python

import re,os
import commands
import subprocess
import signal
from xml.dom import minidom
from collections import defaultdict
import _ped
import parted_devices

if os.environ.get('DEBUG'):
    DEBUG = True
else:
    DEBUG = False

''' algorithm from libparted/unit.c '''
def parse_unit_suffix(suffix):
    if (len(suffix) > 1) and (suffix[1] == 'i'):
        return {
                'k': _ped.UNIT_KIBIBYTE,
                'm': _ped.UNIT_MEBIBYTE,
                'g': _ped.UNIT_GIBIBYTE,
                't': _ped.UNIT_TEBIBYTE,
                }[suffix[0].lower()]
    elif (len(suffix) > 0):
        return {
                's': _ped.UNIT_SECTOR,
                'b': _ped.UNIT_BYTE,
                'k': _ped.UNIT_KILOBYTE,
                'm': _ped.UNIT_MEGABYTE,
                'g': _ped.UNIT_GIGABYTE,
                't': _ped.UNIT_TERABYTE,
                'c': _ped.UNIT_CYLINDER,
                '%': _ped.UNIT_PERCENT,
                }[suffix[0].lower()]

''' type name for display '''
def get_type_name(type):
    try: return {
            _ped.PARTITION_NORMAL       :'primary',
            _ped.PARTITION_LOGICAL      :'logical',
            _ped.PARTITION_EXTENDED     :'extended',
            _ped.PARTITION_FREESPACE    :' ',
            _ped.PARTITION_METADATA     :' ',
            _ped.PARTITION_PROTECTED    :' ',
            }[type]
    except KeyError:
        return ' '

''' get_flags from rhel5 anaconda partedUtils.py
 algorithm the same as parted in parted.c partition_print_flags() '''
def get_flags (part):
    string=""
    if not part.is_active ():
        return string
    flag = _ped.partition_flag_next (0)
    first = 1
    while flag:
        if part.get_flag (flag):
            if first:
                first = 0
                string = string + _ped.partition_flag_get_name(flag) 
            else:
                string = string + ","+ _ped.partition_flag_get_name(flag) 
        flag = _ped.partition_flag_next (flag)
    return string

'''make nested dict'''
def makehash():
    return defaultdict(makehash)

class Partition:
    ''' disk partition '''
    unit = ''
    dict = makehash()
    ''' Partition init function '''
    def __init__(self, dev, ds, pt, dfm, sn, partition, num, st, se, sz, tp, fs, fg, fr, fm, mp):

		'''special process linux-swap(v1) use linux-swap for instead '''
		if fs.startswith('linux-swap'):
			fs = 'linux-swap'
		self.dict[dev][dev] = dev
		self.dict[dev]['disk_size'] = ds
		self.dict[dev]['partition_table'] = pt
		''' if modify parittion table this flag is no '''
		self.dict[dev]['from_os'] = dfm
		self.dict[dev]['partition'][partition]['sn']= sn
		self.dict[dev]['partition'][partition]['num']= num
		self.dict[dev]['partition'][partition]['start'] = st
		self.dict[dev]['partition'][partition]['end'] = se
		self.dict[dev]['partition'][partition]['size'] = sz
		self.dict[dev]['partition'][partition]['type'] = tp
		self.dict[dev]['partition'][partition]['filesystem'] = fs
		self.dict[dev]['partition'][partition]['flags'] = fg
		self.dict[dev]['partition'][partition]['from_os'] = fr
		self.dict[dev]['partition'][partition]['format'] = fm
		self.dict[dev]['partition'][partition]['mount_point'] = mp

    @staticmethod
    def get_disks():
        #return parted_devices.get_disks()
	process = subprocess.Popen("../new_partition/parted_devices.py" ,stdout=subprocess.PIPE, shell=True)
	return eval(process.stdout.readline().rstrip())

    @staticmethod
    def check_partitions():
        cmd = "grep [a-zA-Z][0-9] /proc/partitions>/dev/null"
        ret = os.system(cmd)
        return ret

    @staticmethod
    def init_from_os(unit='b'):

        Partition.unit = unit

        if Partition.unit:
            default_unit = parse_unit_suffix(Partition.unit)
        else:
            default_unit = _ped.UNIT_COMPACT
        _ped.unit_set_default(default_unit)

        ''' each disk '''
        for d in Partition.get_disks():
            _device = _ped.device_get(d)

            dev = _device.path

            ''' algorithm from parted.c do_print() '''
            disk_start = _device.unit_format (0)
            disk_end = _device.unit_format_byte (_device.length * _device.sector_size \
                - (default_unit == _ped.UNIT_CHS or \
                default_unit == _ped.UNIT_CYLINDER))

            disk_size = disk_end

            try:
				_disk = _ped.Disk(_device)
            except Exception, message:
				m_p = re.compile('.*: unrecognised disk label')
				dev = _device.path

				''' process partition table is empty '''
				if m_p.match(str(message)):
					Partition(dev,					# device name
							disk_size,              # disk size
							'',                    # partition_table
							'yes',                  # disk from os
							'',                    # partition sequence number
							'',                    # partition
							'',                    # num
							'',                   	# start
							'',                    # end
							'',                    # size
							'',                    # type
							'',                    # file system
							'',                   	# flags
							'yes',                  # partition from os
							'',                    # whether format
							''                     # mount point
							)

				continue

            ''' free partition number '''
            fn = 0
            ''' partition sequence number '''
            sn = 0

            partition_table = _disk.type.name
            ''' each partition '''
            _part = _disk.next_partition()

            if DEBUG:
                print dev
                print partition_table
                print disk_start
                print disk_size
            while _part:

                ''' skip _part.type = 9
                 I can't find type 9 from pyparted '''
                if _part.type == 9:
                    _part = _disk.next_partition(_part)
                    continue

                p_g = _part.geom

                ''' algorithm get from parted.c do_print() '''
                start = _device.unit_format(p_g.start)
                end = _device.unit_format_byte((p_g.end + 1) * _device.sector_size - 1)
                size = _device.unit_format(p_g.length)

                if _part.num == -1:
                    fn += 1
                    num = fn
                    partition = "free %s" %fn
                else:
                    num = _part.num
                    partition = _part.get_path()

                type_name = get_type_name(_part.type)

                try:
                    fs = _part.fs_type.name
                except:
                    fs = ' '
                ''' may be PARTITION_METADATA mustn't appear at free space,
                 but may be pyparted have bug, awareness last partition as
                 metadata. But parted is not. '''
                if  _part.type & _ped.PARTITION_FREESPACE or \
                        _part.type & _ped.PARTITION_METADATA:
                    fs = 'Free Space'

                flags = get_flags(_part)

                sn += 1

                if DEBUG:
                    print "sn:%s, num:%s, %s, start:%s, end:%s, size:%s,\
type_name:%s, fs: %s, flags: %s" %(sn, num, partition, start, end, size, \
                    type_name, fs, flags)

                ''' create a Partition instance from hardware info '''
                Partition(dev ,                     # device name
                           disk_size,               # disk size
                           partition_table,         # partition table
                           'yes',                   # disk from os
                           sn,                      # partition sequence number
                           partition,               # partition
                           num,                     # num
                           start,                   # start
                           end,                     # end
                           size,                    # size
                           type_name,               # type
                           fs,                      # file system
                           flags,                   # flags
                           'yes',                   # partition from os
                           '',                     # whether format
                           ''                      # mount point
                           )

                ''' next partition '''
                _part = _disk.next_partition(_part)

    @staticmethod
    def init_from_xml(node):
        ''' create Partition instances from xml node '''
        Partition.unit = node.attributes['unit'].value.encode('ascii')
        if DEBUG:
            print Partition.unit
        for d in node.childNodes:
            if d.nodeType == node.ELEMENT_NODE and d.nodeName == 'disk':
                dev = d.attributes['disk'].value.encode('ascii')
                disk_size = d.attributes['disk_size'].value.encode('ascii')
                partition_table = d.attributes['partition_table'].value.encode('ascii')
                disk_from_os = d.attributes['from_os'].value.encode('ascii')

                if DEBUG:
                    print dev
                    print partition_table
                    print disk_size
                    print disk_from_os

                for p in d.getElementsByTagName('partition'):
                    sn = p.attributes['sn'].value.encode('ascii')
                    partition = p.attributes['partition'].value.encode('ascii')
                    num = p.attributes['num'].value.encode('ascii')
                    start = p.attributes['start'].value.encode('ascii')
                    end = p.attributes['end'].value.encode('ascii')
                    size = p.attributes['size'].value.encode('ascii')
                    type_name = p.attributes['type'].value.encode('ascii')
                    fs = p.attributes['filesystem'].value.encode('ascii')
                    flags = p.attributes['flags'].value.encode('ascii')
                    from_os = p.attributes['from_os'].value.encode('ascii')
                    fmt = p.attributes['format'].value.encode('ascii')
                    mp = p.attributes['mount_point'].value.encode('ascii')

                    if DEBUG:
                        print "sn:%s, num:%s, %s, start:%s, end:%s, size:%s,\
type_name:%s, fs: %s, flags: %s" %(sn, num, partition, start, end, size, \
                    type_name, fs, flags)

                    ''' initialise from xml '''
                    Partition(dev ,               # device name
                           disk_size,             # disk size
                           partition_table,       # partition table
                           disk_from_os,			# disk from os
                           sn,                    # partition sequence number
                           partition,             # partition
                           num,                   # num
                           start,                 # start
                           end,                   # end
                           size,                  # size
                           type_name,             # type
                           fs,                    # file system
                           flags,                 # flags
                           from_os,               # from os
                           fmt,                   # whether format
                           mp                     # mount point
                           )


    @staticmethod
    def to_xml(doc, p_node):
        ''' write all Partition instance into xml
doc - xml document instance
p_node - xml node (parent node)'''
        pts = doc.createElement('disks')
        unit_attr = doc.createAttribute('unit')
        unit_attr.value = Partition.unit
        pts.setAttributeNode(unit_attr)

        create_node(Partition.dict, doc, pts)
        p_node.appendChild(pts)
        p_node.appendChild(pts)

    @staticmethod
    def get_disk_attr(disk, attr):
        try:
            return Partition.dict[disk][attr]
        except:
            return None

    @staticmethod
    def get_partition_attr(partition, attr):
       disk =  re.sub("\d+$", "", partition)
       try:
           return Partition.dict[disk]['partition'][partition][attr]
       except:
           return None

    @staticmethod
    def get_disk_from_partition(partition):
		for d in Partition.dict:
			for p in Partition.dict[d]['partition']:
				if p == partition:
					return d

    @staticmethod
    def get_raid_devices():
        rd = []
        for key in Partition.dict:
            for k in Partition.dict[key]['partition'].keys():
                #if Partition.dict[key]['partition'][k]['flags'] == 'raid':
                if re.search('raid',Partition.dict[key]['partition'][k]['flags']):
                    rd.append(k)

        return rd

''' Fix Me:
 make recursive is better '''
def create_node(d, doc, pts):
    for key in d.keys():
        dt = doc.createElement('disk')
        dt.setAttribute('disk', key)
        for k in d[key].keys():
            if k == 'partition':
                pt = doc.createElement('partitions')
                for l in d[key]['partition'].keys():
                    lt = doc.createElement('partition')
                    lt.setAttribute('partition', str(l))
                    for m in d[key]['partition'][l].keys():
                        '''if mount_point etc. is ' ', then can't set attribute,
                        so I manual add space
                        '''
                        attr = "%s " %d[key]['partition'][l][m]
                        if attr.rstrip():
                            attr = attr.rstrip()
                        lt.setAttribute(m, attr)
                    pt.appendChild(lt)
                dt.appendChild(pt)
            elif k == 'disk_size':
                dt.setAttribute(str(k), d[key]['disk_size'])
            elif k == 'partition_table':
                dt.setAttribute(str(k), d[key]['partition_table'])
            elif k == 'from_os':
                dt.setAttribute(str(k), d[key]['from_os'])
        pts.appendChild(dt)

def pretty(d, indent=0):
    for key, value in d.iteritems():
        print '\t' * indent + str(key)
        if isinstance(value, dict):
            pretty(value, indent+1)
        else:
            print '\t' * (indent+1) + str(value)

def delete_all_partitions(d, k):
    d[k]['from_os'] = 'no'
    for key in d[k]['partition'].keys():
        d[k]['partition'].pop(key, None)

    _device = _ped.device_get(k)

    ''' Get end size maybe use disk_new_fresh function is better.'''
    ''' TODO: must use a function to convert to compact unit.'''
    if d[k]['partition_table'] == 'msdos':
		''' algorithm from libparted/labels/dos.c add_startend_metadata()'''
		end = "%ss" %(_device.length - 1)
    elif d[k]['partition_table'] == 'gpt':
		''' algorithm from libparted/labels/gpt.c gpt_alloc()'''
		GPT_DEFAULT_PARTITION_ENTRY_ARRAY_SIZE=int(16384)
		end = "%ss" %(_device.length-2 - \
				GPT_DEFAULT_PARTITION_ENTRY_ARRAY_SIZE / _device.sector_size)
    else:
		end = d[k]['disk_size']


    partition = 'free 1'
    sz = d[k]['disk_size']
    d[k]['partition'][partition]['sn']= 1
    d[k]['partition'][partition]['num']= 1
    d[k]['partition'][partition]['start'] = '0'
    d[k]['partition'][partition]['end'] = end
    d[k]['partition'][partition]['size'] = sz
    d[k]['partition'][partition]['type'] = ''
    d[k]['partition'][partition]['filesystem'] = ''
    d[k]['partition'][partition]['flags'] = ''
    d[k]['partition'][partition]['from_os'] = ''
    d[k]['partition'][partition]['format'] = ''
    d[k]['partition'][partition]['mount_point'] = ''

def delete_one_partition(dict, disk, partition):
    for key in dict[disk]['partition'].keys():
        if key == partition:
            del dict[disk]['partition'][key]

def sort_partition(dict, disk, p):
    f_d = {}
    p_d = {}

    for k in dict[disk]["partition"]:
        if k.startswith("free"):
            f_d[k] = int(dict[disk]["partition"][k]["num"])
        else:
			p_d[k] = int(dict[disk]["partition"][k]["num"])

    if p == "free":
        return (sorted(f_d, key = lambda key: f_d[key]))
    elif p == "partition":
		return (sorted(p_d, key = lambda key: p_d[key]))
    else:
        return []

def sort_partitions(dict, disk, p="all"):
    if dict[disk]['partition_table'] == '':
		return []

    if p == "all":
        ''' k, t: partitions '''
        return (sorted(((k) for k in dict[disk]["partition"]), key=lambda \
                t: int(dict[disk]["partition"][t]["sn"])))
    elif p == "free":
        return sort_partition(dict, disk, "free")
    elif p == "partition":
        return sort_partition(dict, disk, "partition")

def write_to_xml_file():
    xmldoc = minidom.Document()
    root = xmldoc.createElement('install')
    xmldoc.appendChild(root)
    Partition.to_xml(xmldoc, root)
    xml_string =xmldoc.toprettyxml(indent=" ")
    print xml_string

def init_from_xml():
    xmldoc = minidom.parse('install.xml')
    root = xmldoc.firstChild
    for n in root.childNodes:
        if n.nodeType == n.ELEMENT_NODE:
            if n.nodeName == 'disks':
                Partition.init_from_xml(n)

if __name__ == "__main__":
    Partition.init_from_os(unit='s')
    #print Partition.get_disk_from_partition('/dev/sdc7')
#    get_flags('/dev/sdc1')
    #for d in Partition.get_raid_devices():
    #    print d
    #print Partition.dict.keys()
    #pretty(Partition.dict)
    #print Partition.unit
    #write_to_xml_file()
    #disks = Partition.dict.keys()
    #disks.sort()
    #for key in disks:
	#	partitions = sort_partitions(Partition.dict, key, 'partition')
	#	for p in partitions:
	#		for k in Partition.dict[key]['partition'].keys():
	#			if k == p:
	#				print "%s, %s, %s" %(key, k, Partition.dict[key]['partition'][k]['filesystem'])
    #init_from_xml()
    #pretty(Partition.dict)
    #delete_all_partitions(Partition.dict, "/dev/sda")
    #delete_all_partitions(Partition.dict, "/dev/sdf")
    #delete_one_partition(Partition.dict, "/dev/sdb", "free 1")
    #pretty(Partition.dict)
    #print sort_partitions(Partition.dict, "/dev/sda", "all")
    #print "+"*100
    #print sort_partitions(Partition.dict, "/dev/sda", "free")
    #print "+"*100
    #print sort_partitions(Partition.dict, "/dev/sda", "partition")
    #print Partition.dict
    #print Partition.get_size("/dev/sda")
    #print Partition.get_disk_attr("/dev/sda", "disk_size")
    #print Partition.get_disk_attr("/dev/sda", "partition_table")
    #print Partition.get_partition_attr("/dev/sda1", "size")