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 argparse
  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 action_add_cdrom(args):
 138     attrs = {key: vars(args)[key] for key in ('cdrom_id', 'description', 'mirror_id', 'uri_suffix')}
 139
 140     cdrom = database.CDROM(**attrs)
 141     database.session.add(cdrom)
 142     database.session.commit()
 143
 144     load_image(cdrom)
 145
 146 def action_add_mirror(args):
 147     attrs = {key: vars(args)[key] for key in ('mirror_id', 'uri_prefix')}
 148
 149     mirror = database.Mirror(**attrs)
 150     database.session.add(mirror)
 151     database.session.commit()
 152
 153 def action_update(args):
 154     if not args.names:
 155         images = database.CDROM.query().all()
 156     else:
 157         images = [database.CDROM.query().get(arg) for arg in args.names]
 158
 159     for cdrom in images:
 160         if cdrom is not None:
 161             load_image(cdrom)
 162
 163 def action_reap(args):
 164     reap_images()
 165
 166 def main():
 167     database.connect()
 168     database.session.begin()
 169
 170     parser = argparse.ArgumentParser(description='Perform actions on the CD-ROM images in the database')
 171     subparsers = parser.add_subparsers(help='Action to perform')
 172
 173     add_parser = subparsers.add_parser('add', help='Add new image to database')
 174     add_parser.set_defaults(func=lambda args: add_parser.print_help())
 175     add_subparsers = add_parser.add_subparsers()
 176
 177     add_cdrom_parser = add_subparsers.add_parser('cdrom')
 178     add_cdrom_parser.add_argument('cdrom_id')
 179     add_cdrom_parser.add_argument('description')
 180     add_cdrom_parser.add_argument('mirror_id')
 181     add_cdrom_parser.add_argument('uri_suffix')
 182     add_cdrom_parser.set_defaults(func=action_add_cdrom)
 183
 184     add_mirror_parser = add_subparsers.add_parser('mirror')
 185     add_mirror_parser.add_argument('mirror_id')
 186     add_mirror_parser.add_argument('uri_prefix')
 187     add_mirror_parser.set_defaults(func=action_add_mirror)
 188
 189     update_parser = subparsers.add_parser('update', help='Update images in database')
 190     update_parser.add_argument('names', nargs='*', metavar='name', help='Shortnames of images to update')
 191     update_parser.set_defaults(func=action_update)
 192
 193     reap_parser = subparsers.add_parser('reap', help='Reap old images in database')
 194     reap_parser.set_defaults(func=action_reap)
 195
 196     args = parser.parse_args()
 197     if 'func' in args:
 198         args.func(args)
 199     else:
 200         parser.print_help()
 201
 202
 203 if __name__ == '__main__':
 204     main()