Updating version for precise
[invirt/packages/invirt-images.git] / 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()