invirt-images

   1 #!/usr/bin/env python3
   2
   3 import os
   4 import subprocess
   5 import random
   6 import string
   7 import tempfile
   8 import urllib.request
   9 import math
  10 import optparse
  11
  12 from invirt import database
  13
  14
  15 def lvcreate(name, size):
  16     subprocess.run(['lvcreate', '-L', size, '-n', name, 'xenvg'],
  17                    check_output=True, encoding='utf-8', check=True)
  18
  19 def lvrename(dest, src):
  20     subprocess.run(['lvchange', '-an', f'xenvg/{src}'],
  21                    check_output=True, encoding='utf-8', check=True)
  22
  23     subprocess.run(['lvrename', f'xenvg/{src}', f'xenvg/{dest}'],
  24                    check_output=True, encoding='utf-8', check=True)
  25
  26     subprocess.run(['lvchange', '-ay', f'xenvg/{dest}'],
  27                    check_output=True, encoding='utf-8', check=True)
  28
  29 def lv_random(func, pattern, *args):
  30     """
  31     Run some LVM-related command with a random string in the LV name.
  32
  33     func takes an LV name plus whatever's in *args and returns the
  34     return code of some LVM command, such as lvcreate or lvrename
  35
  36     pattern must contain one '{}' pattern, which will be replaced
  37     by a 6-character random string.
  38
  39     the script will attempt to re-run itself if the error code
  40     indicates that the destination already exists
  41     """
  42
  43     # Keep trying until it works
  44     while True:
  45         letters = (random.choice(string.ascii_letters) for _ in range(6))
  46         rand_string = ''.join(letters)
  47
  48         name = pattern.format(rand_string)
  49
  50         try:
  51             func(name, *args)
  52         except subprocess.CalledProcessError as e:
  53             # 5 is the return code if the destination already exists
  54             if e.returncode != 5:
  55                 raise
  56         else:
  57             return name
  58
  59 def fetch_image(cdrom):
  60     """
  61     Download a cdrom from a URI, shelling out to rsync if appropriate
  62     and otherwise trying to use urllib
  63     """
  64
  65     full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
  66     temp_file = tempfile.mkstemp()[1]
  67     print(f'Fetching image {cdrom.cdrom_id} from {full_uri} to {temp_file}')
  68     try:
  69         if full_uri.startswith('rsync://'):
  70             subprocess.run(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
  71                            check_output=True, encoding='utf-8', check=True)
  72         else:
  73             # I'm not going to look for errors here, because I bet it'll
  74             # throw its own exceptions
  75             urllib.request.urlretrieve(full_uri, temp_file)
  76         return temp_file
  77     except:
  78         os.remove(temp_file)
  79         raise
  80
  81 def copy_file(src, dest):
  82     """
  83     Copy a file from one location to another using dd
  84     """
  85
  86     subprocess.run(['dd', f'if={src}', f'of={dest}', 'bs=1M'],
  87                    check_output=True, encoding='utf-8', check=True)
  88
  89 def load_image(cdrom):
  90     """
  91     Update a cdrom image by downloading the latest version,
  92     transferring it into an LV, moving the old LV out of the way and
  93     the new LV into place
  94     """
  95
  96     if cdrom.mirror_id is None:
  97         return
  98
  99     temp_file = fetch_image(cdrom)
 100
 101     st_size = os.stat(temp_file).st_size
 102
 103     assert st_size > 0, 'CD-ROM image size is 0'
 104
 105     megabytes = math.ceil((float(st_size) / (1024 * 1024)))
 106     cdrom_size = f'{megabytes}M'
 107
 108     try:
 109         new_lv = lv_random(lvcreate, f'image-new_{cdrom.cdrom_id}' + '_{}', cdrom_size)
 110         copy_file(temp_file, f'/dev/xenvg/{new_lv}')
 111         lv_random(lvrename, f'image-old_{cdrom.cdrom_id}' + '_{}', 'image_{cdrom.cdrom_id}')
 112         lv_random(lvrename, f'image_{cdrom.cdrom_id}', new_lv)
 113         reap_images()
 114     finally:
 115         os.remove(temp_file)
 116
 117 def reap_images():
 118     """
 119     Remove stale cdrom images that are no longer in use
 120
 121     load_image doesn't attempt to remove the old image because it
 122     might still be in use. reap_images attempts to delete any LVs
 123     starting with 'image-old_', but ignores errors, in case they're
 124     still being used.
 125     """
 126
 127     lvm_list = subprocess.run(['lvs', '-o', 'lv_name', '--noheadings'],
 128                               check_output=True, encoding='utf-8', check=True)
 129
 130     for lv in (s.strip() for s in lvm_list.stdout.readlines()):
 131         if lv.startswith('image-old_'):
 132             subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
 133             subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
 134             subprocess.run(['lvchange', '-a', 'ey', f'/dev/xenvg/{lv}'])
 135             subprocess.run(['lvremove', '--force', f'/dev/xenvg/{lv}'])
 136
 137 def main():
 138     database.connect()
 139     database.session.begin()
 140
 141     usage = """%prog [options] --add [--cdrom] cdrom_id description mirror_id uri_suffix
 142        %prog [options] --add --mirror mirror_id uri_prefix
 143
 144        %prog [options] --update [short_name1 [short_name2 ...]]
 145        %prog [options] --reap"""
 146
 147     parser = optparse.OptionParser(usage=usage)
 148     parser.set_defaults(verbosity=0,
 149                         item='cdrom')
 150
 151     parser.add_option('-a', '--add', action='store_const',
 152                       dest='action', const='add',
 153                       help='Add a new item to the database')
 154
 155     parser.add_option('-u', '--update', action='store_const',
 156                       dest='action', const='update',
 157                       help='Update all cdrom images in the database with the latest version')
 158     parser.add_option('-r', '--reap', action='store_const',
 159                       dest='action', const='reap',
 160                       help='Reap stale cdrom images that are no longer in use')
 161
 162     a_group = optparse.OptionGroup(parser, 'Adding new items')
 163     a_group.add_option('-c', '--cdrom', action='store_const',
 164                        dest='item', const='cdrom',
 165                        help='Add a new cdrom to the database')
 166     a_group.add_option('-m', '--mirror', action='store_const',
 167                        dest='item', const='mirror',
 168                        help='Add a new mirror to the database')
 169     parser.add_option_group(a_group)
 170
 171     v_group = optparse.OptionGroup(parser, "Verbosity levels")
 172     v_group.add_option("-q", "--quiet", action='store_const',
 173                        dest='verbosity', const=0,
 174                        help='Show no output from commands this script runs (default)')
 175     v_group.add_option("-v", "--verbose", action='store_const',
 176                        dest='verbosity', const=1,
 177                        help='Show only errors from commands this script runs')
 178     v_group.add_option("--noisy", action='store_const',
 179                        dest='verbosity', const=2,
 180                        help='Show all output from commands this script runs')
 181     parser.add_option_group(v_group)
 182
 183     (options, args) = parser.parse_args()
 184     if options.action is None:
 185         print(parser.format_help())
 186     elif options.action == 'add':
 187         if options.item == 'cdrom':
 188             attrs = dict(list(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
 189                                   args)))
 190             cdrom = database.CDROM(**attrs)
 191             database.session.add(cdrom)
 192             database.session.commit()
 193
 194             load_image(cdrom)
 195
 196         elif options.item == 'mirror':
 197             attrs = dict(list(zip(('mirror_id', 'uri_prefix'),
 198                                   args)))
 199             mirror = database.Mirror(**attrs)
 200             database.session.add(mirror)
 201             database.session.commit()
 202     elif options.action == 'update':
 203         if len(args) > 0:
 204             images = [database.CDROM.query().get(arg) for arg in args]
 205         else:
 206             images = database.CDROM.query().all()
 207         for cdrom in images:
 208             if cdrom is not None:
 209                 load_image(cdrom)
 210     elif options.action == 'reap':
 211         reap_images()
 212
 213 if __name__ == '__main__':
 214     main()