Fully switch over to Python 3
[invirt/packages/invirt-images.git] / invirt-images
index 2175ead..eeda870 100755 (executable)
-#!/usr/bin/python
+#!/usr/bin/env python3
 
-from invirt import database
 import os
-import sys
 import subprocess
 import random
 import string
 import tempfile
-import urllib
+import urllib.request
 import math
-import optparse as op
-
-class InvirtImageException(Exception):
-    pass
+import optparse
 
-# 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
+from invirt import database
 
-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
+    subprocess.run(['lvcreate', '-L', size, '-n', name, 'xenvg'],
+                   check_output=True, encoding='utf-8', check=True)
 
 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
+    subprocess.run(['lvchange', '-an', f'xenvg/{src}'],
+                   check_output=True, encoding='utf-8', check=True)
+
+    subprocess.run(['lvrename', f'xenvg/{src}', f'xenvg/{dest}'],
+                   check_output=True, encoding='utf-8', check=True)
+
+    subprocess.run(['lvchange', '-ay', f'xenvg/{dest}'],
+                   check_output=True, encoding='utf-8', check=True)
 
 def lv_random(func, pattern, *args):
     """
-    Run some LVM-related command, optionally with a random string in
-    the LV name.
-    
+    Run some LVM-related command 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
+
+    pattern must contain one '{}' pattern, which will be replaced
+    by a 6-character random string.
+
+    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)
+        letters = (random.choice(string.ascii_letters) for _ in range(6))
+        rand_string = ''.join(letters)
 
-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)
+        name = pattern.format(rand_string)
 
-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)
+        try:
+            func(name, *args)
+        except subprocess.CalledProcessError as e:
+            # 5 is the return code if the destination already exists
+            if e.returncode != 5:
+                raise
+        else:
+            return name
 
 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)
+    print(f'Fetching image {cdrom.cdrom_id} from {full_uri} to {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
+            subprocess.run(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
+                           check_output=True, encoding='utf-8', check=True)
         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)
+            urllib.request.urlretrieve(full_uri, temp_file)
         return temp_file
     except:
-        os.unlink(temp_file)
+        os.remove(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)
+
+    subprocess.run(['dd', f'if={src}', f'of={dest}', 'bs=1M'],
+                   check_output=True, encoding='utf-8', check=True)
 
 def load_image(cdrom):
     """
@@ -140,80 +92,74 @@ def load_image(cdrom):
     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
-    try:
-        temp_file = fetch_image(cdrom)
-    except InvirtImageException, e:
-        print >>sys.stderr, 'ERROR: %s.  Skipping.' % e
-        return
+
+    temp_file = fetch_image(cdrom)
+
+    st_size = os.stat(temp_file).st_size
+
+    assert st_size > 0, 'CD-ROM image size is 0'
+
+    megabytes = math.ceil((float(st_size) / (1024 * 1024)))
+    cdrom_size = f'{megabytes}M'
 
     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)
+        new_lv = lv_random(lvcreate, f'image-new_{cdrom.cdrom_id}' + '_{}', cdrom_size)
+        copy_file(temp_file, f'/dev/xenvg/{new_lv}')
+        lv_random(lvrename, f'image-old_{cdrom.cdrom_id}' + '_{}', 'image_{cdrom.cdrom_id}')
+        lv_random(lvrename, f'image_{cdrom.cdrom_id}', new_lv)
         reap_images()
     finally:
-        os.unlink(temp_file)
+        os.remove(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()):
+
+    lvm_list = subprocess.run(['lvs', '-o', 'lv_name', '--noheadings'],
+                              check_output=True, encoding='utf-8', check=True)
+
+    for lv in (s.strip() for s in lvm_list.stdout.readlines()):
         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())
+            subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
+            subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
+            subprocess.run(['lvchange', '-a', 'ey', f'/dev/xenvg/{lv}'])
+            subprocess.run(['lvremove', '--force', f'/dev/xenvg/{lv}'])
 
 def main():
-    global verbosity
-    
     database.connect()
-    
+    database.session.begin()
+
     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 = optparse.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 = optparse.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')
@@ -221,8 +167,8 @@ def main():
                        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 = optparse.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)')
@@ -233,27 +179,26 @@ def main():
                        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()
+        print(parser.format_help())
     elif options.action == 'add':
         if options.item == 'cdrom':
-            attrs = dict(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
-                             args))
+            attrs = dict(list(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
+                                  args)))
             cdrom = database.CDROM(**attrs)
-            database.session.save(cdrom)
-            database.session.flush()
-            
+            database.session.add(cdrom)
+            database.session.commit()
+
             load_image(cdrom)
-        
+
         elif options.item == 'mirror':
-            attrs = dict(zip(('mirror_id', 'uri_prefix'),
-                             args))
+            attrs = dict(list(zip(('mirror_id', 'uri_prefix'),
+                                  args)))
             mirror = database.Mirror(**attrs)
-            database.session.save(mirror)
-            database.session.flush()
+            database.session.add(mirror)
+            database.session.commit()
     elif options.action == 'update':
         if len(args) > 0:
             images = [database.CDROM.query().get(arg) for arg in args]