invirt-images

   1 #!/usr/bin/python
   2
   3 from invirt import database
   4 import os
   5 import sys
   6 import subprocess
   7 import random
   8 import string
   9 import tempfile
  10 import urllib
  11 import math
  12 import optparse as op
  13
  14 class InvirtImageException(Exception):
  15     pass
  16
  17 # verbosity = 0 means no output from the actual commands
  18 # verbosity = 1 means only errors from the actual commands
  19 # verbosity = 2 means all output from the actual commands
  20 verbosity = 0
  21
  22 def getOutput():
  23     global verbosity
  24     return {
  25         'stdout': subprocess.PIPE if verbosity < 2 else None,
  26         'stderr': subprocess.PIPE if verbosity < 1 else None
  27         }
  28
  29 def lvcreate(name, size):
  30     lvc = subprocess.Popen(['lvcreate', '-L', size, '-n', name, 'xenvg'],
  31                            stderr=subprocess.PIPE,
  32                            stdout=getOutput()['stdout'])
  33     if not lvc.wait():
  34         return 0
  35     stderr = lvc.stderr.read()
  36     if 'already exists in volume group' in stderr:
  37         return 5
  38     else:
  39         if verbosity > 0:
  40             print stderr
  41         return 6
  42
  43 def lvrename(dest, src):
  44     lvr = subprocess.Popen(['lvrename', 'xenvg', src, dest],
  45                            stderr=subprocess.PIPE,
  46                            stdout=getOutput()['stdout'])
  47     ret = lvr.wait()
  48     if not ret:
  49         return 0
  50     stderr = lvr.stderr.read()
  51     if 'not found in volume group' in stderr:
  52         return 0
  53     else:
  54         if verbosity > 0:
  55             print stderr
  56         return ret
  57
  58 def lv_random(func, pattern, *args):
  59     """
  60     Run some LVM-related command, optionally with a random string in
  61     the LV name.
  62
  63     func takes an LV name plus whatever's in *args and returns the
  64     return code of some LVM command, such as lvcreate or lvrename
  65
  66     pattern can contain at most one '%s' pattern, which will be
  67     replaced by a 6-character random string.
  68
  69     If pattern contains a '%s', the script will attempt to re-run
  70     itself if the error code indicates that the destination already
  71     exists
  72     """
  73     # Keep trying until it works
  74     while True:
  75         rand_string = ''.join(random.choice(string.ascii_letters) \
  76                                   for i in xrange(6))
  77         if '%s' in pattern:
  78             name = pattern % rand_string
  79         else:
  80             name = pattern
  81         ret = func(name, *args)
  82         if ret == 0:
  83             return name
  84         # 5 is the return code if the destination already exists
  85         elif '%s' not in pattern or ret != 5:
  86             raise InvirtImageException, 'E: Error running %s with args %s' % (func.__name__, args)
  87
  88 def lvcreate_random(pattern, size):
  89     """
  90     Creates an LV, optionally with a random string in the name.
  91
  92     Call with a string formatting pattern with a single '%s' to use as
  93     a pattern for the name of the new LV.
  94     """
  95     return lv_random(lvcreate, pattern, size)
  96
  97 def lvrename_random(src, pattern):
  98     """
  99     Rename an LV to a new name with a random string incorporated.
 100
 101     Call with a string formatting pattern with a single '%s' to use as
 102     a pattern for the name of the new LV
 103     """
 104     return lv_random(lvrename, pattern, src)
 105
 106 def fetch_image(cdrom):
 107     """
 108     Download a cdrom from a URI, shelling out to rsync if appropriate
 109     and otherwise trying to use urllib
 110     """
 111     full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
 112     temp_file = tempfile.mkstemp()[1]
 113     if verbosity > 0:
 114         print >>sys.stderr, "Fetching image %s from %s to %s" % (cdrom.cdrom_id, full_uri, temp_file)
 115     try:
 116         if full_uri.startswith('rsync://'):
 117             if subprocess.call(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
 118                                **getOutput()):
 119                 raise InvirtImageException, "E: Unable to download '%s'" % full_uri
 120         else:
 121             # I'm not going to look for errors here, because I bet it'll
 122             # throw its own exceptions
 123             urllib.urlretrieve(full_uri, temp_file)
 124         return temp_file
 125     except:
 126         os.unlink(temp_file)
 127         raise
 128
 129 def copy_file(src, dest):
 130     """
 131     Copy a file from one location to another using dd
 132     """
 133     if subprocess.call(['dd', 'if=%s' % src, 'of=%s' % dest, 'bs=1M'],
 134                        **getOutput()):
 135         raise InvirtImageException, 'E: Unable to transfer %s into %s' % (src, dest)
 136
 137 def load_image(cdrom):
 138     """
 139     Update a cdrom image by downloading the latest version,
 140     transferring it into an LV, moving the old LV out of the way and
 141     the new LV into place
 142     """
 143     if cdrom.mirror_id is None:
 144         return
 145     temp_file = fetch_image(cdrom)
 146     try:
 147         st_size = os.stat(temp_file).st_size
 148         if not st_size:
 149             print >>sys.stderr, "Failed to fetch %s" % cdrom.cdrom_id
 150             return
 151         cdrom_size = '%sM' % math.ceil((float(st_size) / (1024 * 1024)))
 152         new_lv = lvcreate_random('image-new_%s_%%s' % cdrom.cdrom_id, cdrom_size)
 153         copy_file(temp_file, '/dev/xenvg/%s' % new_lv)
 154         lvrename_random('image_%s' % cdrom.cdrom_id, 'image-old_%s_%%s' % cdrom.cdrom_id)
 155         lvrename_random(new_lv, 'image_%s' % cdrom.cdrom_id)
 156         reap_images()
 157     finally:
 158         os.unlink(temp_file)
 159
 160 def reap_images():
 161     """
 162     Remove stale cdrom images that are no longer in use
 163
 164     load_image doesn't attempt to remove the old image because it
 165     might still be in use. reap_images attempts to delete any LVs
 166     starting with 'image-old_', but ignores errors, in case they're
 167     still being used.
 168     """
 169     lvm_list = subprocess.Popen(['lvs', '-o', 'lv_name', '--noheadings'],
 170                                stdout=subprocess.PIPE,
 171                                stdin=subprocess.PIPE)
 172     lvm_list.wait()
 173
 174     for lv in map(str.strip, lvm_list.stdout.read().splitlines()):
 175         if lv.startswith('image-old_'):
 176             subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
 177                             **getOutput())
 178             subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
 179                             **getOutput())
 180             subprocess.call(['lvchange', '-a', 'ey', '/dev/xenvg/%s' % lv],
 181                             **getOutput())
 182             subprocess.call(['lvremove', '--force', '/dev/xenvg/%s' % lv],
 183                             **getOutput())
 184
 185 def main():
 186     global verbosity
 187
 188     database.connect()
 189
 190     usage = """%prog [options] --add [--cdrom] cdrom_id description mirror_id uri_suffix
 191        %prog [options] --add --mirror mirror_id uri_prefix
 192
 193        %prog [options] --update [short_name1 [short_name2 ...]]
 194        %prog [options] --reap"""
 195
 196     parser = op.OptionParser(usage=usage)
 197     parser.set_defaults(verbosity=0,
 198                         item='cdrom')
 199
 200     parser.add_option('-a', '--add', action='store_const',
 201                       dest='action', const='add',
 202                       help='Add a new item to the database')
 203
 204     parser.add_option('-u', '--update', action='store_const',
 205                       dest='action', const='update',
 206                       help='Update all cdrom images in the database with the latest version')
 207     parser.add_option('-r', '--reap', action='store_const',
 208                       dest='action', const='reap',
 209                       help='Reap stale cdrom images that are no longer in use')
 210
 211     a_group = op.OptionGroup(parser, 'Adding new items')
 212     a_group.add_option('-c', '--cdrom', action='store_const',
 213                        dest='item', const='cdrom',
 214                        help='Add a new cdrom to the database')
 215     a_group.add_option('-m', '--mirror', action='store_const',
 216                        dest='item', const='mirror',
 217                        help='Add a new mirror to the database')
 218     parser.add_option_group(a_group)
 219
 220     v_group = op.OptionGroup(parser, "Verbosity levels")
 221     v_group.add_option("-q", "--quiet", action='store_const',
 222                        dest='verbosity', const=0,
 223                        help='Show no output from commands this script runs (default)')
 224     v_group.add_option("-v", "--verbose", action='store_const',
 225                        dest='verbosity', const=1,
 226                        help='Show only errors from commands this script runs')
 227     v_group.add_option("--noisy", action='store_const',
 228                        dest='verbosity', const=2,
 229                        help='Show all output from commands this script runs')
 230     parser.add_option_group(v_group)
 231
 232     (options, args) = parser.parse_args()
 233     verbosity = options.verbosity
 234     if options.action is None:
 235         print parser.format_help()
 236     elif options.action == 'add':
 237         if options.item == 'cdrom':
 238             attrs = dict(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
 239                              args))
 240             cdrom = database.CDROM(**attrs)
 241             database.session.save(cdrom)
 242             database.session.flush()
 243
 244             load_image(cdrom)
 245
 246         elif options.item == 'mirror':
 247             attrs = dict(zip(('mirror_id', 'uri_prefix'),
 248                              args))
 249             mirror = database.Mirror(**attrs)
 250             database.session.save(mirror)
 251             database.session.flush()
 252     elif options.action == 'update':
 253         if len(args) > 0:
 254             images = [database.CDROM.query().get(arg) for arg in args]
 255         else:
 256             images = database.CDROM.query().all()
 257         for cdrom in images:
 258             if cdrom is not None:
 259                 load_image(cdrom)
 260     elif options.action == 'reap':
 261         reap_images()
 262
 263 if __name__ == '__main__':
 264     main()