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