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     try:
 146         temp_file = fetch_image(cdrom)
 147     except InvirtImageException, e:
 148         print >>sys.stderr, 'ERROR: %s.  Skipping.' % e
 149         return
 150
 151     try:
 152         st_size = os.stat(temp_file).st_size
 153         if not st_size:
 154             print >>sys.stderr, "Failed to fetch %s" % cdrom.cdrom_id
 155             return
 156         cdrom_size = '%sM' % math.ceil((float(st_size) / (1024 * 1024)))
 157         new_lv = lvcreate_random('image-new_%s_%%s' % cdrom.cdrom_id, cdrom_size)
 158         copy_file(temp_file, '/dev/xenvg/%s' % new_lv)
 159         lvrename_random('image_%s' % cdrom.cdrom_id, 'image-old_%s_%%s' % cdrom.cdrom_id)
 160         lvrename_random(new_lv, 'image_%s' % cdrom.cdrom_id)
 161         reap_images()
 162     finally:
 163         os.unlink(temp_file)
 164
 165 def reap_images():
 166     """
 167     Remove stale cdrom images that are no longer in use
 168
 169     load_image doesn't attempt to remove the old image because it
 170     might still be in use. reap_images attempts to delete any LVs
 171     starting with 'image-old_', but ignores errors, in case they're
 172     still being used.
 173     """
 174     lvm_list = subprocess.Popen(['lvs', '-o', 'lv_name', '--noheadings'],
 175                                stdout=subprocess.PIPE,
 176                                stdin=subprocess.PIPE)
 177     lvm_list.wait()
 178
 179     for lv in map(str.strip, lvm_list.stdout.read().splitlines()):
 180         if lv.startswith('image-old_'):
 181             subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
 182                             **getOutput())
 183             subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
 184                             **getOutput())
 185             subprocess.call(['lvchange', '-a', 'ey', '/dev/xenvg/%s' % lv],
 186                             **getOutput())
 187             subprocess.call(['lvremove', '--force', '/dev/xenvg/%s' % lv],
 188                             **getOutput())
 189
 190 def main():
 191     global verbosity
 192
 193     database.connect()
 194     database.session.begin()
 195
 196     usage = """%prog [options] --add [--cdrom] cdrom_id description mirror_id uri_suffix
 197        %prog [options] --add --mirror mirror_id uri_prefix
 198
 199        %prog [options] --update [short_name1 [short_name2 ...]]
 200        %prog [options] --reap"""
 201
 202     parser = op.OptionParser(usage=usage)
 203     parser.set_defaults(verbosity=0,
 204                         item='cdrom')
 205
 206     parser.add_option('-a', '--add', action='store_const',
 207                       dest='action', const='add',
 208                       help='Add a new item to the database')
 209
 210     parser.add_option('-u', '--update', action='store_const',
 211                       dest='action', const='update',
 212                       help='Update all cdrom images in the database with the latest version')
 213     parser.add_option('-r', '--reap', action='store_const',
 214                       dest='action', const='reap',
 215                       help='Reap stale cdrom images that are no longer in use')
 216
 217     a_group = op.OptionGroup(parser, 'Adding new items')
 218     a_group.add_option('-c', '--cdrom', action='store_const',
 219                        dest='item', const='cdrom',
 220                        help='Add a new cdrom to the database')
 221     a_group.add_option('-m', '--mirror', action='store_const',
 222                        dest='item', const='mirror',
 223                        help='Add a new mirror to the database')
 224     parser.add_option_group(a_group)
 225
 226     v_group = op.OptionGroup(parser, "Verbosity levels")
 227     v_group.add_option("-q", "--quiet", action='store_const',
 228                        dest='verbosity', const=0,
 229                        help='Show no output from commands this script runs (default)')
 230     v_group.add_option("-v", "--verbose", action='store_const',
 231                        dest='verbosity', const=1,
 232                        help='Show only errors from commands this script runs')
 233     v_group.add_option("--noisy", action='store_const',
 234                        dest='verbosity', const=2,
 235                        help='Show all output from commands this script runs')
 236     parser.add_option_group(v_group)
 237
 238     (options, args) = parser.parse_args()
 239     verbosity = options.verbosity
 240     if options.action is None:
 241         print parser.format_help()
 242     elif options.action == 'add':
 243         if options.item == 'cdrom':
 244             attrs = dict(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
 245                              args))
 246             cdrom = database.CDROM(**attrs)
 247             database.session.add(cdrom)
 248             database.session.commit()
 249
 250             load_image(cdrom)
 251
 252         elif options.item == 'mirror':
 253             attrs = dict(zip(('mirror_id', 'uri_prefix'),
 254                              args))
 255             mirror = database.Mirror(**attrs)
 256             database.session.add(mirror)
 257             database.session.commit()
 258     elif options.action == 'update':
 259         if len(args) > 0:
 260             images = [database.CDROM.query().get(arg) for arg in args]
 261         else:
 262             images = database.CDROM.query().all()
 263         for cdrom in images:
 264             if cdrom is not None:
 265                 load_image(cdrom)
 266     elif options.action == 'reap':
 267         reap_images()
 268
 269 if __name__ == '__main__':
 270     main()