"""
Xen Migrate
Migrate XenServer to Open Source Xen
2009 Mark Pace -- Jolokia Networks
pace@jolokianetworks.com
GPL License
USE THIS SOFTWARE AT YOUR OWN RISK!
"""

import os
import subprocess
import sys

def docmd(cmd):
    """
    run a command and return the communicate PIPE
    """
    execute=subprocess.Popen([cmd],shell=True,stdout=subprocess.PIPE)
    return execute.communicate()[0]

def exportvm(vmname,lvdev,destfile):
    """
    export lvdev to dest
    """
    # we'll need to handle difference block sizes at some point
    blocksize=1024*1024
    vmuuid=getvmuuid(vmname)
    vmstatus=getvmstatus(vmuuid)
    if vmstatus=='running':
        cmd='xe vm-shutdown -u root uuid='+vmuuid
        docmd(cmd)
    vmstatus=getvmstatus(vmuuid)
    if vmstatus=='halted':
        print '\nActivating Volume:'
        cmd='lvchange -v -ay '+lvdev
        lvchange=docmd(cmd)
        try:
            source=open(lvdev,'rb')
            if not os.path.exists(destfile):
                dest=open(destfile,'wb')
                print '\nRW notification every: '+str((blocksize*1024)/(2**30))+'GB' # 2**30=GB
                sys.stdout.write('Exporting: ')
                write=0
                while True:
                    write=write+1
                    data=source.read(blocksize)
                    if write%(blocksize/1024)==0:
                        sys.stdout.write(str(write/1024)+'GBr')
                    if len(data)==0:
                        break #EOF
                    dest.write(data)
                    if write%(blocksize/1024)==0:
                        sys.stdout.write('w ')
                    sys.stdout.flush()
            else:
                print 'ERROR: destination file '+destfile+' exists.'
        finally:
            try:
                print
                source.close()
                dest.close()
                print 'Successful export'
            finally:
                print '\nDeactivating Volume:'
                cmd='lvchange -v -an '+lvdev
                docmd(cmd)
                print 'Restart VM with:'
                print 'xe vm-startup -u root uuid='+vmuuid
    else:
        print 'ERROR: vm status:',vmstatus,'vm needs to be halted to migrate'

def importvm(lvdest,sourcefile,vgdest,lvsize):
    """
    import a raw vmfile into a logical volume
    """
    blocksize=1024*1024
    lvexists=0
    lvvgs=getlvdevlist()
    for lvvg in lvvgs:
        if lvdest==lvvg[0]:
            print 'ERROR: lv '+lvdest+' exists cannot import'
            lvexists=1
    if not lvexists:
        cmd='lvcreate -v -n '+lvdest+' -L '+lvsize+'G '+vgdest
        print '\nCreating Logical Volume:'
        docmd(cmd)
        try:
            source=open(sourcefile,'rb')
            destlv=vgdest+'/'+lvdest
            dest=open(destlv,'wb')
            print '\nRW notification every: '+str((blocksize*1024)/(2**30))+'GB' # 2**30=GB
            sys.stdout.write('Importing: ')
            write=0
            while True:
                write=write+1
                data=source.read(blocksize)
                if write%(blocksize/1024)==0:
                    sys.stdout.write(str(write/1024)+'GBr')
                if len(data)==0:
                    break # EOF
                dest.write(data)
                if write%(blocksize/1024)==0:
                    sys.stdout.write('w ')
                sys.stdout.flush()
        finally:
            try:
                print
                source.close()
                dest.close()
                print 'Successful import'
            finally:
                print
    else:
        print 'ERROR: logical volume '+lvdest+' exists'

def getlvdevlist():
    """
    get logical volume and volume group list and return it
    """
    lvvgs=[]
    sep=','
    cmd='lvs --separator \''+sep+'\''
    vgdevs=docmd(cmd).split('\n')
    del vgdevs[0]
    del vgdevs[-1]
    for vgdev in vgdevs:
        lv=vgdev.split(sep)[0][2:]
        vg=vgdev.split(sep)[1]
        size=vgdev.split(sep)[3][:-1]
        lvvgs.append([lv,vg,size])
    return lvvgs

def getlvdevxen(vmdiskuuid):
    """
    take the vmdisk uuid and return the logical volume device name
    """
    lvvgs=getlvdevlist()
    for lvvg in lvvgs:
        if vmdiskuuid in lvvg[0]:
            lvdev='/dev/'+lvvg[1]+'/'+lvvg[0]
            return lvdev,lvvg[2]

def getvmdiskuuid(vmuuid):
    """
    get the vmdisk uuids from the vmuuid
    return disk uuids in list
    """
    diskuuid=[]
    cmd='xe vbd-list vm-uuid='+vmuuid
    response=docmd(cmd).split('vdi-uuid ( RO): ')
    del response[0]
    for index,uuid in enumerate(response):
        curuuid=uuid.split('\n')[0]
        if curuuid!='<not in database>':
            partid=uuid.split('\n')[2].split(': ')[1]
            diskuuid.append([curuuid,partid])
    return diskuuid

def getvmstatus(vmuuid):
    cmd='xe vm-list uuid='+vmuuid
    response=docmd(cmd).split('power-state ( RO): ')[1].split('\n')[0]
    return response

def getvmuuid(vmname):
    """
    get the vmuuid from the name-label of a vm
    return uuid
    """
    try:
        cmd='xe vm-list name-label=\''+vmname+'\''
        uuid=docmd(cmd).split(':')[1].split(' ')[1][:-1]
        return uuid
    except IndexError:
        return 'vm not found'

def reftoraw(refdir,rawfile):
    """
    take the ref directory of an xva file and create a raw importable file
    """
    blocksize=1024*1024
    refdirlist=os.listdir(refdir)
    numfiles=int(refdirlist[-2])
    if os.path.isdir(refdir):
        if not os.path.exists(rawfile):
            try:
                filenum=0
                print '\nRW notification every: '+str((blocksize*1024)/(2**30))+'GB' # 2**30=GB
                dest=open(rawfile,'wr')
                sys.stdout.write('\nConverting: ')
                while filenum<=numfiles:
                    if (filenum+1)%(blocksize/1024)==0:
                        sys.stdout.write(str((filenum+1)/1024)+'GBr')
                    filename=str(filenum)
                    while len(filename)<8:
                        filename='0'+filename
                    if os.path.exists(refdir+filename):
                        source=open(refdir+filename,'rb')
                        while True:
                            data=source.read(blocksize)
                            if len(data)==0:
                                source.close()
                                break # EOF
                            dest.write(data)
                    else:
                        dest.seek(blocksize,1)
                    if (filenum+1)%(blocksize/1024)==0:
                        sys.stdout.write('w ')
                    sys.stdout.flush()
                    filenum+=1
            finally:
                print
                try:
                    dest.close()
                    source.close()
                    print 'Successful convert'
                finally:
                    print
        else:
            print 'ERROR: rawfile '+rawfile+' exists'
    else:
        print 'ERROR: refdir '+refdir+' does not exist'
                    
##
## Main Program
##

if __name__=='__main__':
    from optparse import OptionParser
    parser=OptionParser(usage='%prog [options] vm-name|lvdest|rawfile')
    parser.add_option('-d','--disk',action='store_true',dest='disk',help='display vm disk uuids',default=False)
    parser.add_option('-l','--lvdev',action='store_true',dest='lvdev',help='display vm logical volume devices',default=False)
    parser.add_option('-x','--export',action='store',type='string',dest='export',metavar='FILE',help='export to FILE')
    parser.add_option('-i','--import',action='store',type='string',dest='doimport',metavar='FILE',help='import from FILE')
    parser.add_option('-c','--convert',action='store',type='string',dest='convert',metavar='DIR',help='convert xva ref dir to importable rawfile')
    (opts,args)=parser.parse_args()
    print '\nxenmigrate\n'
    if len(args)<1:
        parser.print_help()
        sys.exit(1)
    if opts.disk or opts.lvdev or opts.export:
        vmname=args[0]
        vmuuid=getvmuuid(vmname)
        print 'vm name-label    :',vmname
        print 'vm uuid          :',vmuuid
        vmdiskuuids=getvmdiskuuid(vmuuid)
        for vmdiskuuid in vmdiskuuids:
            print 'vm disk uuid     :',vmdiskuuid[0]
            print 'vm disk partid   :',vmdiskuuid[1]
            if opts.lvdev:
                lvdev=getlvdevxen(vmdiskuuid[0])
                if lvdev is not None:
                    lvdevname,lvsize=getlvdevxen(vmdiskuuid[0])
                    print 'vm disk dev name :',lvdevname
                    print 'vm disk size     :',lvsize+'GB'
                else:
                    print 'vm disk dev name : not found in mounted storage repositories'
    if opts.export and opts.doimport:
        print 'error: export and import cannot be run at the same time'
    elif opts.export:
        vmdiskuuids=getvmdiskuuid(vmuuid)
        for vmdiskuuid in vmdiskuuids:
            lvdev,lvsize=getlvdevxen(vmdiskuuid[0])
            exportname=opts.export+'_'+vmdiskuuid[1]+'_'+lvsize
            print 'exporting        :',lvdev
            print 'to               :',exportname
            if lvdev:
                exportvm(vmname,lvdev,exportname)
    elif opts.doimport:
        lvsize=opts.doimport.split('_')[-1]
        lvpartid=opts.doimport.split('_')[-2]
        lvdest=opts.doimport.split('_')[0].split('/')[-1]
        print 'lv dest           :',lvdest
        print 'vm partid         :',lvpartid
        print 'vm size           :',lvsize+'GB'
        importvm(lvdest,opts.doimport,args[0],lvsize)
    elif opts.convert:
        reftoraw(opts.convert,args[0])