cp-virsh.py

#!/usr/bin/python
#
# cp_virsh.py -pete 2016-03-15
#
#

'''
usage: cp-virsh.py [-h] [--remove] [--quiet] [--password PASSWORD]
                   [--memory MEMORY] [--cpus CPUS] [--add-disk]
                   [--disk-size DISK_SIZE] [--remove-vol]
                   domain

Provisions new libvirt VMs based on an existing domain/image. This probably
only works when run with python 2.7 on linux

positional arguments:
  domain                name of existing domain

optional arguments:
  -h, --help            show this help message and exit
  --remove, -r          removes the domain and its disk image
  --quiet, -q           run quitely
  --password PASSWORD, -p PASSWORD
                        root password on guest
  --memory MEMORY, -m MEMORY
                        memory in gigabits
  --cpus CPUS, -c CPUS  number of virtual CPUs
  --add-disk            adds a storage device to a domain
  --disk-size DISK_SIZE
                        size in GB of disk for --add-disk
  --remove-vol          clean up the extra storage matching the domain pattern
'''

import libvirt
import pexpect
import subprocess
import sys
import re
import argparse
import os
import time
from xml.dom import minidom



TEMP_XML        = "/tmp/temp.xml"
DISK_TEMPLATE   = "/var/lib/libvirt/images/{}.qcow2"
NAME_PATTERN    = re.compile("([6a-z]*)(\d*)")
STORAGE_PATTERN = re.compile("([6a-z]*)(\d*?)-(\d*?)\.qcow2")
STORAGE_POOL    = "default"
HOSTNAME_SCRIPT = "/root/chg_hostname.sh"
HOST_BRIDGE     = 'br0'
GUEST_USERNAME  = 'root'
BOOT_DELAY      = 20
CONSOLE_DELAY   = 20
SHELL_PS_MATCH  = '%s@%s'
MEM_PATTERN1    = '/<memory/ s/.*/<memory unit="KiB">{}<\/memory>/'
MEM_PATTERN2    = '/<currentMemory/ s/.*/<currentMemory unit="KiB">{}<\/currentMemory>/'
CPU_PATTERN1    = '/<vcpu / s/.*/<vcpu placement="static">{}<\/vcpu>/'
MIN_MEM         = 1
MAX_MEM         = 12
MIN_CPU         = 1
MAX_CPU         = 10
VOL_XML         = """
<volume type='file'>
  <name>{name}.qcow2</name>
  <key>{path}</key>
  <source>
  </source>
  <capacity unit='G'>{size}</capacity>
  <allocation unit='bytes'>0</allocation>
  <target>
    <path>{path}</path>
    <format type='qcow2'/>
    <permissions>
      <mode>0600</mode>
      <owner>108</owner>
      <group>116</group>
    </permissions>
    <compat>1.1</compat>
    <features>
      <lazy_refcounts/>
    </features>
  </target>
</volume>
"""



# Run a command on the system
def run_command(command, verify=False, quiet=False):

    if verify:
        print " ".join(command)
        print "you're about to run this command... continue?"
        c = raw_input("(y/n): ")
        if c != "y": return False

    if not quiet: print "EXEC: " + " ".join(command)
    f = subprocess.call(command)

    if f > 0:
        print "EXEC FAILED!!!"
        return False
    return True



# Connect to the local libvirt daemon
def get_connection():
    return libvirt.open("qemu:///system")



# Save all vm names in a list
def get_domains(conn):
    return [e.name() for e in conn.listAllDomains()]



# Figure out next hostname to use as new_name
def get_new_hostname(domains, template_name, name_pattern):

    # Find template naming pattern
    h = name_pattern.match(template_name).groups()[0]

    # Get the number portion of the names
    nums = []
    for e in domains:
        m = name_pattern.match(e)
        if m and m.groups()[0] == h:
            nums.append(int(m.groups()[1]))

    if len(nums) < 1:
        print >>sys.stderr, "FAIL: no existing domains found matching pattern "+template_name
        raise Exception
        #return None

    # Increment hostname by 1 or find gap
    nums.sort()
    cur = nums[0] + 1
    for e in nums[1:]:
        if e != cur: break
        cur += 1

    return "{0}{1:02d}".format(h, cur)




# Take a copy of the template and save it to a temp file
def get_xml_doc(conn, template_name, temp_xml):
    xml = minidom.parseString(conn.lookupByName(template_name).XMLDesc())
    print >>open(temp_xml, "wb"), xml.toxml()



# Returns gigabit size in kibibit
def parse_mem(gigabit):
    try:
        if int(gigabit) < MIN_MEM or int(gigabit) > MAX_MEM:
            raise Exception("incorrect memory size")
        return int(pow(2, 20) * int(gigabit))
    except Exception, e:
        raise Exception("parse_mem FAIL: Must specify a digit between {} and {} (ERR: {})".format(MIN_MEM, MAX_MEM,e))



# Verifies CPU string is valid
def parse_cpu(cpus):
    try:
        if int(cpus) < MIN_CPU or int(cpus) > MAX_CPU:
            raise Exception("incorrect cpu count")
        return int(cpus)
    except Exception, e:
        raise Exception("parse_cpu FAIL: Must specify a digit between {} and {} (ERR: {})".format(MIN_CPU, MAX_CPU,e))



def modify_xml_doc(temp_xml, template_name, new_name, memory, cpus, quiet):
    # Remove uuid
    run_command(['sed', '-i', '/uuid/d', temp_xml], quiet=quiet)
    # Change name
    run_command(['sed', '-i', 's/%s/%s/' % (template_name, new_name), temp_xml], quiet=quiet)
    # Remove interface
    run_command(['sed', '-i', '/interface/,/interface/d', temp_xml], quiet=quiet)
    # Update memory
    run_command(['sed', '-i', MEM_PATTERN1.format(memory), temp_xml], quiet=quiet)
    run_command(['sed', '-i', MEM_PATTERN2.format(memory), temp_xml], quiet=quiet)
    # Update CPUs
    run_command(['sed', '-i', CPU_PATTERN1.format(cpus), temp_xml], quiet=quiet)



def create_vm(temp_xml, disk_template, new_name, quiet):
    # Clone disk
    run_command(['virsh', 'vol-clone', disk_template, '%s.qcow2' % new_name], quiet=quiet)
    # Create the VM
    run_command(['virsh', 'define', '%s' % temp_xml], quiet=quiet)



def start_vm(domain, quiet):
    run_command(['virsh', 'start', domain], quiet=quiet)



def reboot_vm(domain, quiet):
    run_command(['virsh', 'reboot', domain], quiet=quiet)



def add_interface(domain, quiet):
    run_command(['virsh', 'attach-interface', domain, 'bridge', HOST_BRIDGE, '--persistent'], quiet=quiet)



# Run the chg_hostname script on the new VM
def run_chg_hostname(new_hostname, template_name, username, password, quiet=False):

    if not quiet: print "Connecting to %s's console..." % new_hostname
    try:
        #child = pexpect.spawn('virsh console %s' % new_hostname, logfile=open("test.log", "wb"))
        child = pexpect.spawn('virsh console %s' % new_hostname)
        time.sleep(CONSOLE_DELAY)

        child.sendline()
        match = '%s login:' % template_name
        child.expect(match)

        child.sendline(username)
        match = 'Password: '
        child.expect(match)

        child.sendline(password)
        match = SHELL_PS_MATCH % (username, template_name)
        child.expect(match)

        child.sendline('sh "%s" "%s"' % (HOSTNAME_SCRIPT, new_hostname))
        match = SHELL_PS_MATCH % (username, template_name)
        child.expect(match)

        child.sendline('logout')
        child.sendcontrol(']')

    except pexpect.TIMEOUT, e:
        print >>sys.stderr, "ERROR!"
        print >>sys.stderr, e.get_trace()
        print >>sys.stderr, "Timeout on console session... here is what we see:"
        print >>sys.stderr, child.buffer
        print >>sys.stderr, "We were expecting:"
        print >>sys.stderr, match
        child.sendcontrol(']')
        child.terminate()



# Adds a volume and attaches it to domain
def add_disk(domain, size, quiet):

    conn        = get_connection()
    pool        = conn.storagePoolLookupByName(STORAGE_POOL)
    vol_num     = get_vol_name(domain, pool, STORAGE_PATTERN)

    vol_alpha   = chr(vol_num + ord("a"))
    vol_name    = "{}-{:02d}".format(domain, vol_num)
    vol_target  = "vd{}".format(vol_alpha)
    vol_path    = DISK_TEMPLATE.format(vol_name)
    vol_xml     = VOL_XML.format(
                    name=vol_name,
                    size=size,
                    path=vol_path)

    print "Creating new storage volume here: {}".format(vol_path)
    pool.createXML(vol_xml, 0)
    run_command(["virsh",
        "attach-disk",
        domain,
        vol_path,
        vol_target,
        "--targetbus",
        "virtio",
        "--driver",
        "qemu",
        "--subdriver",
        "qcow2",
        "--live",#"--config"],
        "--persistent"],
    verify=not quiet, quiet=quiet)



# Gets the next disk name for a domain
def get_vol_name(domain, pool, name_pattern):

    vol_list = [e for e in pool.listVolumes()]
    nums = []
    for e in vol_list:
        m = name_pattern.match(e)
        if m and m.groups()[0]+m.groups()[1] == domain:
            nums.append(int(m.groups()[2]))

    if len(nums) == 0: cur = 1
    else:
        nums.sort()
        cur = 1
        for e in nums:
            if e != cur: break
            cur += 1

    return cur



# Returns a list of storage volumes that match name_pattern
def find_vol_list(domain, pool, name_pattern):
    vol_list = [e for e in pool.listVolumes()]
    result = []
    for e in vol_list:
        m = name_pattern.match(e)
        if m and m.groups()[0]+m.groups()[1] == domain:
            result.append(e)

    return result



def remove_domain(domain, quiet):

    disk_image = DISK_TEMPLATE.format(domain)

    run_command(["virsh", "shutdown", domain], verify=not quiet, quiet=quiet)
    run_command(["virsh", "undefine", domain], verify=not quiet, quiet=quiet)
    run_command(["virsh", "vol-delete", "--pool", STORAGE_POOL,
        os.path.basename(disk_image)], verify=not quiet, quiet=quiet)
    remove_vol(domain, quiet)



def remove_vol(domain, quiet):
    conn        = get_connection()
    pool        = conn.storagePoolLookupByName(STORAGE_POOL)

    for e in find_vol_list(domain, pool, STORAGE_PATTERN):
        vol = DISK_TEMPLATE.replace(".qcow2","").format(e)
        run_command(["virsh", "vol-delete", "--pool", STORAGE_POOL, vol],
            verify=not quiet, quiet=quiet)



def create_domain(domain, password, memory, cpus, quiet):

    disk_image  = DISK_TEMPLATE     .format(domain)
    conn        = get_connection    ()
    domains     = get_domains       (conn)
    new_name    = get_new_hostname  (domains, domain, NAME_PATTERN)

    if not new_name: return

    try:
        memory = parse_mem(memory)
        cpus   = parse_cpu(cpus)
    except Exception, e:
        print e
        return

    get_xml_doc(conn, domain, TEMP_XML)
    conn.close()
    modify_xml_doc(TEMP_XML, domain, new_name, memory, cpus, quiet)
    create_vm(TEMP_XML, disk_image, new_name, quiet)

    start_vm(new_name, quiet)
    time.sleep(BOOT_DELAY)
    run_chg_hostname(new_name, domain, GUEST_USERNAME, password)
    reboot_vm(new_name, quiet)
    add_interface(new_name, quiet)

    print "Finished trying to create "+ new_name

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Provisions new libvirt VMs '+
        'based on an existing domain/image. This probably only works when run '+
        'with python 2.7 on linux')

    parser.add_argument("domain", help="name of existing domain")
    parser.add_argument("--remove", "-r", action="store_true", help="removes "+
        "the domain and its disk image")
    parser.add_argument("--quiet", "-q", default=False, action="store_true", help="run quitely")
    parser.add_argument("--password", '-p', help="root password on guest")
    parser.add_argument("--memory", '-m', default=1, help="memory in gigabits")
    parser.add_argument("--cpus", '-c', default=1, help="number of virtual CPUs")
    parser.add_argument("--add-disk", default=False, action="store_true", help="adds a storage device to a domain")
    parser.add_argument("--disk-size", default="8", help="size in GB of disk for --add-disk")
    parser.add_argument("--remove-vol", action="store_true", help="clean up the extra storage matching the domain pattern")

    args = parser.parse_args()

    if args.remove:
        remove_domain(args.domain, args.quiet)
    elif args.add_disk:
        add_disk(args.domain, args.disk_size, args.quiet)
    elif args.remove_vol:
        remove_vol(args.domain, args.quiet)
    else:
        if not args.password:
            import getpass
            args.password = getpass.getpass('root password on %s:' % args.domain)

        create_domain(args.domain, args.password, args.memory, args.cpus, args.quiet)