source: trunk/packages/invirt-images/invirt-images @ 2974

#!/usr/bin/python

from invirt import database
import os
import sys
import subprocess
import random
import string
import tempfile
import urllib
import math
import optparse as op

class InvirtImageException(Exception):
    pass

# verbosity = 0 means no output from the actual commands
# verbosity = 1 means only errors from the actual commands
# verbosity = 2 means all output from the actual commands
verbosity = 0

def getOutput():
    global verbosity
    return {
        'stdout': subprocess.PIPE if verbosity < 2 else None,
        'stderr': subprocess.PIPE if verbosity < 1 else None
        }

def lvcreate(name, size):
    lvc = subprocess.Popen(['lvcreate', '-L', size, '-n', name, 'xenvg'],
                           stderr=subprocess.PIPE,
                           stdout=getOutput()['stdout'])
    if not lvc.wait():
        return 0
    stderr = lvc.stderr.read()
    if 'already exists in volume group' in stderr:
        return 5
    else:
        if verbosity > 0:
            print stderr
        return 6

def lvrename(dest, src):
    lvr = subprocess.Popen(['lvrename', 'xenvg', src, dest],
                           stderr=subprocess.PIPE,
                           stdout=getOutput()['stdout'])
    ret = lvr.wait()
    if not ret:
        return 0
    stderr = lvr.stderr.read()
    if 'not found in volume group' in stderr:
        return 0
    else:
        if verbosity > 0:
            print stderr
        return ret

def lv_random(func, pattern, *args):
    """
    Run some LVM-related command, optionally with a random string in
    the LV name.
    
    func takes an LV name plus whatever's in *args and returns the
    return code of some LVM command, such as lvcreate or lvrename
    
    pattern can contain at most one '%s' pattern, which will be
    replaced by a 6-character random string.
    
    If pattern contains a '%s', the script will attempt to re-run
    itself if the error code indicates that the destination already
    exists
    """
    # Keep trying until it works
    while True:
        rand_string = ''.join(random.choice(string.ascii_letters) \
                                  for i in xrange(6))
        if '%s' in pattern:
            name = pattern % rand_string
        else:
            name = pattern
        ret = func(name, *args)
        if ret == 0:
            return name
        # 5 is the return code if the destination already exists
        elif '%s' not in pattern or ret != 5:
            raise InvirtImageException, 'E: Error running %s with args %s' % (func.__name__, args)

def lvcreate_random(pattern, size):
    """
    Creates an LV, optionally with a random string in the name.
    
    Call with a string formatting pattern with a single '%s' to use as
    a pattern for the name of the new LV.
    """
    return lv_random(lvcreate, pattern, size)

def lvrename_random(src, pattern):
    """
    Rename an LV to a new name with a random string incorporated.
    
    Call with a string formatting pattern with a single '%s' to use as
    a pattern for the name of the new LV
    """
    return lv_random(lvrename, pattern, src)

def fetch_image(cdrom):
    """
    Download a cdrom from a URI, shelling out to rsync if appropriate
    and otherwise trying to use urllib
    """
    full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
    temp_file = tempfile.mkstemp()[1]
    if verbosity > 0:
        print >>sys.stderr, "Fetching image %s from %s to %s" % (cdrom.cdrom_id, full_uri, temp_file)
    try:
        if full_uri.startswith('rsync://'):
            if subprocess.call(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
                               **getOutput()):
                raise InvirtImageException, "E: Unable to download '%s'" % full_uri
        else:
            # I'm not going to look for errors here, because I bet it'll
            # throw its own exceptions
            urllib.urlretrieve(full_uri, temp_file)
        return temp_file
    except:
        os.unlink(temp_file)
        raise

def copy_file(src, dest):
    """
    Copy a file from one location to another using dd
    """
    if subprocess.call(['dd', 'if=%s' % src, 'of=%s' % dest, 'bs=1M'],
                       **getOutput()):
        raise InvirtImageException, 'E: Unable to transfer %s into %s' % (src, dest)

def load_image(cdrom):
    """
    Update a cdrom image by downloading the latest version,
    transferring it into an LV, moving the old LV out of the way and
    the new LV into place
    """
    if cdrom.mirror_id is None:
        return
    temp_file = fetch_image(cdrom)
    try:
        st_size = os.stat(temp_file).st_size
        if not st_size:
            print >>sys.stderr, "Failed to fetch %s" % cdrom.cdrom_id
            return
        cdrom_size = '%sM' % math.ceil((float(st_size) / (1024 * 1024)))
        new_lv = lvcreate_random('image-new_%s_%%s' % cdrom.cdrom_id, cdrom_size)
        copy_file(temp_file, '/dev/xenvg/%s' % new_lv)
        lvrename_random('image_%s' % cdrom.cdrom_id, 'image-old_%s_%%s' % cdrom.cdrom_id)
        lvrename_random(new_lv, 'image_%s' % cdrom.cdrom_id)
        reap_images()
    finally:
        os.unlink(temp_file)

def reap_images():
    """
    Remove stale cdrom images that are no longer in use
    
    load_image doesn't attempt to remove the old image because it
    might still be in use. reap_images attempts to delete any LVs
    starting with 'image-old_', but ignores errors, in case they're
    still being used.
    """
    lvm_list = subprocess.Popen(['lvs', '-o', 'lv_name', '--noheadings'],
                               stdout=subprocess.PIPE,
                               stdin=subprocess.PIPE)
    lvm_list.wait()
    
    for lv in map(str.strip, lvm_list.stdout.read().splitlines()):
        if lv.startswith('image-old_'):
            subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
                            **getOutput())
            subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
                            **getOutput())
            subprocess.call(['lvchange', '-a', 'ey', '/dev/xenvg/%s' % lv],
                            **getOutput())
            subprocess.call(['lvremove', '--force', '/dev/xenvg/%s' % lv],
                            **getOutput())

def main():
    global verbosity
    
    database.connect()
    
    usage = """%prog [options] --add [--cdrom] cdrom_id description mirror_id uri_suffix
       %prog [options] --add --mirror mirror_id uri_prefix

       %prog [options] --update [short_name1 [short_name2 ...]]
       %prog [options] --reap"""
    
    parser = op.OptionParser(usage=usage)
    parser.set_defaults(verbosity=0,
                        item='cdrom')
    
    parser.add_option('-a', '--add', action='store_const',
                      dest='action', const='add',
                      help='Add a new item to the database')
    
    parser.add_option('-u', '--update', action='store_const',
                      dest='action', const='update',
                      help='Update all cdrom images in the database with the latest version')
    parser.add_option('-r', '--reap', action='store_const',
                      dest='action', const='reap',
                      help='Reap stale cdrom images that are no longer in use')
    
    a_group = op.OptionGroup(parser, 'Adding new items')
    a_group.add_option('-c', '--cdrom', action='store_const',
                       dest='item', const='cdrom',
                       help='Add a new cdrom to the database')
    a_group.add_option('-m', '--mirror', action='store_const',
                       dest='item', const='mirror',
                       help='Add a new mirror to the database')
    parser.add_option_group(a_group)
    
    v_group = op.OptionGroup(parser, "Verbosity levels")
    v_group.add_option("-q", "--quiet", action='store_const',
                       dest='verbosity', const=0,
                       help='Show no output from commands this script runs (default)')
    v_group.add_option("-v", "--verbose", action='store_const',
                       dest='verbosity', const=1,
                       help='Show only errors from commands this script runs')
    v_group.add_option("--noisy", action='store_const',
                       dest='verbosity', const=2,
                       help='Show all output from commands this script runs')
    parser.add_option_group(v_group)
    
    (options, args) = parser.parse_args()
    verbosity = options.verbosity
    if options.action is None:
        print parser.format_help()
    elif options.action == 'add':
        if options.item == 'cdrom':
            attrs = dict(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
                             args))
            cdrom = database.CDROM(**attrs)
            database.session.save(cdrom)
            database.session.flush()
            
            load_image(cdrom)
        
        elif options.item == 'mirror':
            attrs = dict(zip(('mirror_id', 'uri_prefix'),
                             args))
            mirror = database.Mirror(**attrs)
            database.session.save(mirror)
            database.session.flush()
    elif options.action == 'update':
        if len(args) > 0:
            images = [database.CDROM.query().get(arg) for arg in args]
        else:
            images = database.CDROM.query().all()
        for cdrom in images:
            if cdrom is not None:
                load_image(cdrom)
    elif options.action == 'reap':
        reap_images()

if __name__ == '__main__':
    main()