#! /usr/bin/python

############################################
# Program:   virt-back.py
# Author:    Russell.Ballestrini.net
# Date:      Fri Mar 18 15:56:59 EDT 2011
# License:   Public Domain
############################################

DESCRIPTION = """A backup utility for QEMU, KVM, XEN, and Virtualbox guests.
Virt-back is a python application that uses the libvirt api to safely 
shutdown, gzip, and restart guests.  The backup process logs to syslog
for auditing and virt-back works great with cron for scheduling outages.
Virt-back has been placed in the public domain and 
the latest version may be downloaded here:
https://bitbucket.org/russellballestrini/virt-back
"""

"""The varible doms represents a list of libvirt dom objects.
Use the Domfetcher class to aquire lists of dom objects."""

import libvirt
import tarfile
import syslog
from re import findall
from time import sleep
from datetime import date
from sys import exit
from optparse import OptionParser, OptionGroup
from os import path
from os import remove
from shutil import move, copy2

try:
    from operator import methodcaller
except ImportError:
    def methodcaller(name, *args, **kwargs):
        def caller(obj):
            return getattr(obj, name)(*args, **kwargs)
        return caller

class Domfetcher(object):
    """Abstract libvirt API, supply methods to return dom object lists"""
    def __init__(self, uri = None):
        """Connect to hypervisor uri with read write access"""
        # register logit as the error handler
        libvirt.registerErrorHandler( logit, 'libvirt error' ) 
    
        try:
            self.c = libvirt.open( uri )
        except:
            self.c = None
            logit( "libvirt error", "Failed to open connection to the hypervisor" )

    def get_running_doms(self):
        """Return a list of running dom objects"""
        doms = []
        for id in self.c.listDomainsID(): # loop over the running ids
            dom = self.c.lookupByID( id ) # fetch dom object by id
            if 'Domain-' not in dom.name(): # prevent actions on the Xen hypervisor
                doms.append( dom ) # append dom object to doms
        return doms
    
    def get_shutoff_doms(self):
        """Return a list of all shutoff but defined dom objects"""
        return [self.c.lookupByName( name ) for name in self.c.listDefinedDomains()]
    
    def get_doms_by_names(self, guest_names):
        """Accept a list of guest_names, return a list of related dom objects"""
        doms = []
        for name in guest_names:
            try:
                dom = self.c.lookupByName( name )
                doms.append( dom )
            except libvirt.libvirtError: pass # logit reg'd as libvirt error handler
        return doms

    def get_all_doms(self):
        """Return a list of all dom objects"""
        return self.get_running_doms() + self.get_shutoff_doms()
        

def invoke( doms, method ):
    """Pattern to invoke shutdown, destroy, and start on a list of doms"""
    f = methodcaller( method )
    for dom in doms:
        try:
            logit( method, 'invoking %s on %s' % (method, dom.name()) )
            retcode = f(dom)
            if retcode: # log retcode
                logit( 
                  method,
                  '{0} returned {1} on {2}'.format( method, retcode, dom.name() )
                ) 
        except libvirt.libvirtError:
            pass

def backup( doms ):
    """Accept a list of dom objects, run backup proceedure on each"""
    for dom in doms:
        recreate = dom.isActive()

        if dom.isActive(): # if dom is active, shutdown
            shutdown([dom])
            
        if dom.isActive(): # if dom is active, error
            logit(
              'error', 
              'unable to shutdown or destroy %s and BACKUP FAILED!' % dom.name()
            )
            continue # skip to the next dom

        xml = dom.XMLDesc( 0 )
        xmlfile = path.join( options.backpath, dom.name() + '.xml' )
        f = open( xmlfile, 'w')
        f.write( xml )
        f.close()
        disklist = findall( "<source file='(.*)'/>\n", xml )

        logit( 'backup', 'invoking backup for ' + dom.name() )
        
        for disk_source in disklist:
            disk_file = disk_source.split('/')[-1]
            disk_dest = path.join( options.backpath, disk_file )

            logit(
              'backup', 
              'copying %s to %s for %s' % (disk_source, disk_dest, dom.name())
            )
            copy2( disk_source, disk_dest )
            
        if recreate: # if true, start guest after backup
            create([dom]) # start dom

        ext, tarmode = '.tar.gz', 'w:gz'
        if options.nogzip: 
            ext, tarmode = '.tar', 'w'

        tarfilename = dom.name() + ext 
        if options.tardate: 
            tarfilename = dom.name() + '-' + TODAY + ext 

        tarpath = path.join( options.backpath, tarfilename )
            
        if path.isfile( tarpath ): # if file exists, run rotate
            logit( 'backup', 'rotating backup files for ' + dom.name() )
            rotate( tarpath, options.retention )

        logit( 
          'backup', 
          'archiving files for %s to %s' % (dom.name(), tarpath)
        )
        tar = tarfile.open( tarpath, tarmode )

        logit( 
          'backup', 
          'archiving %s for %s' % (xmlfile, dom.name())
        )
        tar.add( xmlfile ) # add xml to tar
        remove( xmlfile ) # cleanup tmp files

        for disk_source in disklist:
            disk_file = disk_source.split('/')[-1]
            disk_dest = path.join( options.backpath, disk_file )
            logit(
              'backup', 
              'archiving %s for %s' % (disk_dest, dom.name())
            )
            tar.add( disk_dest ) # add img to tar 
            remove( disk_dest ) # cleanup tmp files
            
        tar.close()

        logit( 'backup', 'finished backup for ' + dom.name() )
            
def shutdown( doms, wait=180 ):
    """Accept a list of dom objects, attempt to shutdown the active ones"""
    # get all running guests from list and invoke shutdown   
    invoke( get_all_running( doms ), 'shutdown' )

    """loop until all guests are shut off or destroy 
       all active guests if wait timer is reached."""

    secs = 10
    wait /= secs # divide wait by secs 

    for i in range( 0, wait + 1 ):

        # if all doms are shut off, leave loop
        if check_all_shutoff( doms ): 
            break
        else:
            logit(
              'shutdown', 
              'waited ' + str( i * secs ) + ' seconds for ' + 
              ', '.join( dom.name() for dom in get_all_running( doms ) ) +
              ' to shut off'
            )

        # if the wait time is reached, destroy all active doms
        if i == wait: 
            invoke( get_all_running( doms ) ,'destroy' )
        
        sleep( secs )

def create( doms ):
    """Accept a list of dom objects, attempt to start the inactive ones"""
    # get all shutoff guests from list and invoke create   
    invoke( get_all_shutoff( doms ), 'create' )

def reboot( doms ):
    """Accept a list of dom objects, attempt to shutdown then start"""
    shutdown( doms )
    create( doms )
    
def info( doms ):
    """Accept a list of dom objects, attempt to display info for all"""
    #invoke( doms, 'name' )
    #invoke( doms, 'info')
    if check_all_running( doms ): print "NOTE: All guests are running"
    if check_all_shutoff( doms ): print "NOTE: All guests are shut off"
    
    print ''
    print 'running guests: ' + ', '.join( [ dom.name() for dom in get_all_running( doms ) ] )
    print 'shutoff guests: ' + ', '.join( [ dom.name() for dom in get_all_shutoff( doms ) ] )
    print ''
    print 'DomName'.ljust(16) + 'Memory MB'.rjust(12) + 'vCPUs'.rjust(8) + 'CPUtime ms'.rjust(18)
    print '======================================================'
    for dom in doms:
        name = dom.name()
        rams = str(dom.info()[2]/1024) + '/' + str(dom.info()[1]/1024)
        cpus = str(dom.info()[3])
        time = str(dom.info()[4]/1000000)
        print name.ljust(16) + rams.rjust(12) + cpus.rjust(8) + time.rjust(18)

def check_all_running( doms ):
    """Accept a list of dom objects, check if all guest dom are active"""
    if sum( [dom.isActive() for dom in doms] ) == len( doms ):
        return True
    return False

def check_all_shutoff( doms ):
    """Accept a list of dom objects, check if all guest dom are shut off"""
    if sum( [dom.isActive() for dom in doms] ):
        return False
    return True

def get_all_running( doms ):
    """Accept a list of dom objects, return a list of running dom objects"""
    return [ dom for dom in doms if dom.isActive() ]

def get_all_shutoff( doms ):
    """Accept a list of dom objects, return a list of shutoff dom objects"""
    return [ dom for dom in doms if not dom.isActive() ]

def logit( context, message, quiet = False ):
    """syslog and error handler"""
    if type( message ) is tuple:
        message = message[2] # libvirt message is a tuple

    try: quiet = options.quiet
    except NameError: pass 

    if quiet: pass 
    else: print context + ': ' +  message

    syslog.openlog( 'virt-back', 0, syslog.LOG_LOCAL3 )
    syslog.syslog( message )
    syslog.closelog()

def rotate( target, retention = 3 ):
    """file rotation routine"""
    for i in range( retention-2, 0, -1 ): # count backwards
        old_name = "%s.%s" % ( target, i )
        new_name = "%s.%s" % ( target, i + 1 )
        try: move( old_name, new_name  )
        except IOError: pass
    move( target, target + '.1' )

def getoptions():
    """Fetch cli args, parse and map to python, test sanity"""

    # create an option parcer object
    p = OptionParser()

    p.set_description( DESCRIPTION )
    
    p.add_option( '-q', '--quiet', dest='quiet', action='store_true', default=False,
                  help='prevent output to stdout' )

    p.add_option( '-d', '--date', dest='tardate', action='store_true', default=False,
                  help='append date to tar filename [default: no date]' )

    p.add_option( '-g', '--no-gzip', dest='nogzip', action='store_true', default=False,
                  help='do not gzip the resulting tar file' )

    p.add_option( '-a', '--retention', dest='retention', metavar="amount", default=3,
                  type='int', help="backups to retain [default: 3]" )

    p.add_option( '-p', '--path', dest='backpath', metavar="'PATH'", default='/KVMBACK',
                  help="backup path [default: '/KVMBACK']" )
    
    p.add_option( '-u', '--uri', dest='uri', metavar="'URI'",
                  help="optional hypervisor uri" )

    # Actions for info testing: These options display info/ test a list of guests only.
                  
    g0 = OptionGroup( p, "Actions for info testing", "These options display info or test a list of guests."  )
                  
    g0.add_option( '-i', '--info', dest='info', action='store_true', default=False,
                  help='info/test a list of guests (space delimited dom names)' )
                  
    g0.add_option( '--info-all', dest='infoall', action='store_true', default=False,
                  help='attempt to show info on ALL guests' )

    # WARNING: Dangerous options below, option grouping for scary actions
                  
    g1 = OptionGroup( p, "Actions for a list of dom names", "WARNING:  These options WILL bring down guests!"  )
    
    g1.add_option( '-b', '--backup', dest='backup', action='store_true', default=False,
                  help='backup a list of guests (space delimited dom names)' )
                  
    g1.add_option( '-r', '--reboot', dest='reboot', action='store_true', default=False,
                  help='reboot a list of guests (space delimited dom names)' )
                  
    g1.add_option( '-s', '--shutdown', dest='shutdown', action='store_true', default=False,
                  help='shutdown a list of guests (space delimited dom names)' )
                  
    g1.add_option( '-c', '--create', dest='create', action='store_true', default=False,
                  help='start a list of guests (space delimited dom names)' )
                  
    g2 = OptionGroup( p, "Actions for all doms", "WARNING:  These options WILL bring down ALL guests!"  )

    g2.add_option( '--backup-all', dest='backupall', action='store_true', default=False,
                  help='attempt to shutdown, backup, and start ALL guests' )

    g2.add_option( '--reboot-all', dest='rebootall', action='store_true', default=False,
                  help='attempt to shutdown and then start ALL guests' )
                  
    g2.add_option( '--shutdown-all', dest='shutdownall', action='store_true', default=False,
                  help='attempt to shutdown ALL guests' )

    g2.add_option( '--create-all', dest='createall', action='store_true', default=False,
                  help='attempt to start ALL guests' )

    # attach groups g1 and g2 to the parser p object
    p.add_option_group( g0 )
    p.add_option_group( g1 )
    p.add_option_group( g2 )

    # parse options and args
    options, guest_names = p.parse_args()

    # the actionsum should be 1 to continue, bool math ftw
    actions = [ options.backup, options.reboot, options.shutdown,
      options.create, options.info, options.backupall, options.rebootall, 
      options.shutdownall, options.createall, options.infoall, ]
    actionsum = sum( actions )

    if actionsum == 1:
        return options, guest_names
    else:
        exit("\nYou must have 1 action, no more, no less.\n\nRun 'virt-back --help' for help.\n")

if __name__ == '__main__':
    TODAY = str( date.today() )
    # Get the list of options and list of guest_names from cli
    options, guest_names = getoptions()
    # connect to hypervisor with Domfetcher (read-only)
    domfetcher = Domfetcher( options.uri )  
    
    if options.backup or options.reboot or options.shutdown or options.create or options.info:
        doms = domfetcher.get_doms_by_names( guest_names )
    else: 
        doms = domfetcher.get_all_doms()
        
    if options.backup or options.backupall: backup( doms )
    if options.reboot or options.rebootall: reboot( doms )
    if options.shutdown or options.shutdownall: shutdown( doms )
    if options.create or options.createall: create( doms )
    if options.info or options.infoall: info( doms )


