ansible/library/virt

#!/usr/bin/python 
"""
Virt management features

Copyright 2007, 2012 Red Hat, Inc
Michael DeHaan <mdehaan@redhat.com>
Seth Vidal <skvidal@fedoraproject.org>

This software may be freely redistributed under the terms of the GNU
general public license.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
"""

VIRT_FAILED = 1
VIRT_SUCCESS = 0
VIRT_UNAVAILABLE=2

try:
    import json
except ImportError:
    import simplejson as json


# other modules
import os
import sys
import subprocess
try:
    import libvirt
except ImportError:
    print json.dumps({
        "failed" : True,
        "msg" : "libvirt python module unavailable",
        "rc": VIRT_UNAVAILABLE,
        })   
    sys.exit(1)


import shlex

VIRT_STATE_NAME_MAP = {
   0 : "running",
   1 : "running",
   2 : "running",
   3 : "paused",
   4 : "shutdown",
   5 : "shutdown",
   6 : "crashed"
}

class LibvirtConnection(object):

    def __init__(self):

        cmd = subprocess.Popen("uname -r", shell=True, stdout=subprocess.PIPE,
                                close_fds=True)
        output = cmd.communicate()[0]

        if output.find("xen") != -1:
            conn = libvirt.open(None)
        else:
            conn = libvirt.open("qemu:///system")

        if not conn:
            raise Exception("hypervisor connection failure")

        self.conn = conn

    def find_vm(self, vmid):
        """
        Extra bonus feature: vmid = -1 returns a list of everything
        """
        conn = self.conn

        vms = []

        # this block of code borrowed from virt-manager:
        # get working domain's name
        ids = conn.listDomainsID()
        for id in ids:
            vm = conn.lookupByID(id)
            vms.append(vm)
        # get defined domain
        names = conn.listDefinedDomains()
        for name in names:
            vm = conn.lookupByName(name)
            vms.append(vm)

        if vmid == -1:
            return vms

        for vm in vms:
            if vm.name() == vmid:
                return vm

        raise Exception("virtual machine %s not found" % vmid)

    def shutdown(self, vmid):
        return self.find_vm(vmid).shutdown()

    def pause(self, vmid):
        return self.suspend(self.conn,vmid)

    def unpause(self, vmid):
        return self.resume(self.conn,vmid)

    def suspend(self, vmid):
        return self.find_vm(vmid).suspend()

    def resume(self, vmid):
        return self.find_vm(vmid).resume()

    def create(self, vmid):
        return self.find_vm(vmid).create()

    def destroy(self, vmid):
        return self.find_vm(vmid).destroy()

    def undefine(self, vmid):
        return self.find_vm(vmid).undefine()

    def get_status2(self, vm):
        state = vm.info()[0]
        # print "DEBUG: state: %s" % state
        return VIRT_STATE_NAME_MAP.get(state,"unknown")

    def get_status(self, vmid):
        state = self.find_vm(vmid).info()[0]
        return VIRT_STATE_NAME_MAP.get(state,"unknown")

    def nodeinfo(self):
        return self.conn.getInfo()

    def get_type(self):
        return self.conn.getType()

    def get_maxVcpus(self, vmid):
        vm = self.conn.lookupByName(vmid)
        return vm.maxVcpus()

    def get_maxMemory(self, vmid):
        vm = self.conn.lookupByName(vmid)
        return vm.maxMemory()

    def getFreeMemory(self):
        return self.conn.getFreeMemory()

    def get_autostart(self, vmid):
        vm = self.conn.lookupByName(vmid)
        return vm.autostart()
    
    def set_autostart(self, vmid, val):
        vm = self.conn.lookupByName(vmid)
        return vm.setAutostart(val)
    


class Virt(object):

    def __get_conn(self):
        self.conn = LibvirtConnection()
        return self.conn

    def get_vm(self, vmid):
        self.__get_conn()
        return self.conn.find_vm(vmid)
        
    def state(self):
        vms = self.list_vms()
        state = []
        for vm in vms:
            state_blurb = self.conn.get_status(vm)
            state.append("%s %s" % (vm,state_blurb))
        return state


    def info(self):
        vms = self.list_vms()
        info = dict()
        for vm in vms:
            data = self.conn.find_vm(vm).info()
            # libvirt returns maxMem, memory, and cpuTime as long()'s, which
            # xmlrpclib tries to convert to regular int's during serialization.
            # This throws exceptions, so convert them to strings here and
            # assume the other end of the xmlrpc connection can figure things
            # out or doesn't care.
            info[vm] = {
                "state"     : VIRT_STATE_NAME_MAP.get(data[0],"unknown"),
                "maxMem"    : str(data[1]),
                "memory"    : str(data[2]),
                "nrVirtCpu" : data[3],
                "cpuTime"   : str(data[4]),
            }
            thisvm = self.conn.find_vm(vm)
            if hasattr(thisvm, 'autostart'):
                info[vm]["autostart"] = thisvm.autostart()

        return info

    def nodeinfo(self):
        self.__get_conn()
        info = dict()
        data = self.conn.nodeinfo()
        info = {
            "cpumodel"     : str(data[0]),
            "phymemory"    : str(data[1]),
            "cpus"         : str(data[2]),
            "cpumhz"       : str(data[3]),
            "numanodes"    : str(data[4]),
            "sockets"      : str(data[5]),
            "cpucores"     : str(data[6]),
            "cputhreads"   : str(data[7])
        }
        return info

    def list_vms(self):
        self.conn = self.__get_conn()
        vms = self.conn.find_vm(-1)
        results = []
        for x in vms:
            try:
                results.append(x.name())
            except:
                pass
        return results

    def virttype(self):
        return self.__get_conn().get_type()
        

    def autostart(self, vmid):
        self.conn = self.__get_conn()
        return self.conn.set_autostart(vmid, True)

    def freemem(self):
        self.conn = self.__get_conn()
        return self.conn.getFreeMemory()

    def shutdown(self, vmid):
        """
        Make the machine with the given vmid stop running.
        Whatever that takes.
        """
        self.__get_conn()
        self.conn.shutdown(vmid)
        return 0


    def pause(self, vmid):

        """
        Pause the machine with the given vmid.
        """
        self.__get_conn()
        return self.conn.suspend(vmid)


    def unpause(self, vmid):

        """
        Unpause the machine with the given vmid.
        """

        self.__get_conn()
        return self.conn.resume(vmid)


    def create(self, vmid):

        """
        Start the machine via the given mac address.
        """
        self.__get_conn()
        return self.conn.create(vmid)

    def start(self, vmid):

        """
        Start the machine via the given id/name
        """
        self.__get_conn()
        return self.conn.create(vmid)


    def destroy(self, vmid):

        """
        Pull the virtual power from the virtual domain, giving it virtually no
        time to virtually shut down.
        """
        self.__get_conn()
        return self.conn.destroy(vmid)


    def undefine(self, vmid):

        """
        Stop a domain, and then wipe it from the face of the earth.
        by deleting the disk image and it's configuration file.
        """

        self.__get_conn()
        return self.conn.undefine(vmid)


    def status(self, vmid):

        """
        Return a state suitable for server consumption.  Aka, codes.py values, not XM output.
        """

        self.__get_conn()
        return self.conn.get_status(vmid)


    def get_xml(self, vmid):
        """
        Recieve a Vm id as input
        Return an xml describing vm config returned by a libvirt call
        """
        conn = libvirt.openReadOnly(None)
        if conn == None:
            return (-1,'Failed to open connection to the hypervisor')
        try:
            domV = conn.lookupByName(vmid)
        except:
            return (-1,'Failed to find the main domain')
        return domV.XMLDesc(0)


    def get_maxVcpus(self, vmid):
        """
        Gets the max number of VCPUs on a guest
        """

        self.__get_conn()
        return self.conn.get_maxVcpus(vmid)

    def get_max_memory(self, vmid):
        """
        Gets the max memory on a guest
        """

        self.__get_conn()
        return self.conn.get_MaxMemory(vmid)

def main():
    rc = VIRT_SUCCESS
    vm_commands = ['create','status', 'start', 'stop', 'pause', 'unpause', 
                'shutdown', 'undefine', 'destroy', 'get_xml', 'autostart']
    host_commands = ['freemem', 'list_vms', 'info', 'nodeinfo', 'virttype']
    
    
    msg = """
    virtmodule arguments:
    - queryvm guest <queryable item from libvirt domain api> <more of those>
    - guest commands: (all require a single guest name an argument)
       %s
    - host commands:
       %s
    
    """ % (', '.join(sorted(vm_commands)), ', '.join(sorted(host_commands)))
    
    if len(sys.argv) == 1:
        return VIRT_FAILED, msg

    argfile = sys.argv[1]
    if not os.path.exists(argfile):
        msg = "Argument file not found"
        return VIRT_FAILED, msg

    args = open(argfile, 'r').read()
    items = shlex.split(args)

    if not len(items):
        return VIRT_FAILED, msg

    basecmd = items[0]
    v = Virt()

    if basecmd in ('queryvm'):
        if len(items) < 3:
            msg = "queryvm requires at least 2 arguments: guest itemname [itemname]"
            return VIRT_FAILED, msg
        
        guest = items[1]
        reqs = items[2:]
        res = {}
        for req in reqs:
            vm = v.get_vm(guest)
            if hasattr(vm, req):
                data = getattr(vm,req)()
                res[req] = data
        return rc, res
        
    # any methods that require the guest/vmid as the argument must be listed here
    elif basecmd in vm_commands:
        if len(items) < 2:
            msg = "%s requires 1 argument: guest" % basecmd
        return VIRT_FAILED, msg
                         
        guest = items[1]
        res = getattr(v, basecmd)(guest)
        if type(res) != dict:
            res = { basecmd: res }
        return rc, res
    
    elif hasattr(v, basecmd):
        res = getattr(v, basecmd)()
        return rc, res
        
    else:
        msg = "Command %s not recognized" % basecmd
        rc = VIRT_FAILED
        
    return rc, msg
    
if __name__ == "__main__":
    try:
        rc, result = main()
    except Exception, e:
        rc = 1
        result = str(e)

    if rc != 0: # something went wrong emit the msg
        print json.dumps({
            "failed" : rc,
            "msg" : result,
            "rc": rc,
        })   
        sys.exit(rc)
    else:
        print json.dumps(result)