Restore html_theme = 'default' in docs/conf.py
[snf-image-creator] / image_creator / main.py
1 #!/usr/bin/env python
2
3 # Copyright 2012 GRNET S.A. All rights reserved.
4 #
5 # Redistribution and use in source and binary forms, with or
6 # without modification, are permitted provided that the following
7 # conditions are met:
8 #
9 #   1. Redistributions of source code must retain the above
10 #      copyright notice, this list of conditions and the following
11 #      disclaimer.
12 #
13 #   2. Redistributions in binary form must reproduce the above
14 #      copyright notice, this list of conditions and the following
15 #      disclaimer in the documentation and/or other materials
16 #      provided with the distribution.
17 #
18 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
19 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
20 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
21 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
22 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
25 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
26 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
27 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
28 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 # POSSIBILITY OF SUCH DAMAGE.
30 #
31 # The views and conclusions contained in the software and
32 # documentation are those of the authors and should not be
33 # interpreted as representing official policies, either expressed
34 # or implied, of GRNET S.A.
35
36 from image_creator import __version__ as version
37 from image_creator.disk import Disk
38 from image_creator.util import FatalError, MD5
39 from image_creator.output.cli import SilentOutput, SimpleOutput, \
40     OutputWthProgress
41 from image_creator.os_type import os_cls
42 from image_creator.kamaki_wrapper import Kamaki, ClientError
43 import sys
44 import os
45 import optparse
46 import StringIO
47 import signal
48
49
50 def check_writable_dir(option, opt_str, value, parser):
51     dirname = os.path.dirname(value)
52     name = os.path.basename(value)
53     if dirname and not os.path.isdir(dirname):
54         raise FatalError("`%s' is not an existing directory" % dirname)
55
56     if not name:
57         raise FatalError("`%s' is not a valid file name" % dirname)
58
59     setattr(parser.values, option.dest, value)
60
61
62 def parse_options(input_args):
63     usage = "Usage: %prog [options] <input_media>"
64     parser = optparse.OptionParser(version=version, usage=usage)
65
66     parser.add_option("-o", "--outfile", type="string", dest="outfile",
67                       default=None, action="callback",
68                       callback=check_writable_dir, help="dump image to FILE",
69                       metavar="FILE")
70
71     parser.add_option("-f", "--force", dest="force", default=False,
72                       action="store_true",
73                       help="overwrite output files if they exist")
74
75     parser.add_option("-s", "--silent", dest="silent", default=False,
76                       help="output only errors",
77                       action="store_true")
78
79     parser.add_option("-u", "--upload", dest="upload", type="string",
80                       default=False,
81                       help="upload the image to pithos with name FILENAME",
82                       metavar="FILENAME")
83
84     parser.add_option("-r", "--register", dest="register", type="string",
85                       default=False,
86                       help="register the image with ~okeanos as IMAGENAME",
87                       metavar="IMAGENAME")
88
89     parser.add_option("-m", "--metadata", dest="metadata", default=[],
90                       help="add custom KEY=VALUE metadata to the image",
91                       action="append", metavar="KEY=VALUE")
92
93     parser.add_option("-t", "--token", dest="token", type="string",
94                       default=None, help="use this authentication token when "
95                       "uploading/registering images")
96
97     parser.add_option("--print-sysprep", dest="print_sysprep", default=False,
98                       help="print the enabled and disabled system preparation "
99                       "operations for this input media", action="store_true")
100
101     parser.add_option("--enable-sysprep", dest="enabled_syspreps", default=[],
102                       help="run SYSPREP operation on the input media",
103                       action="append", metavar="SYSPREP")
104
105     parser.add_option("--disable-sysprep", dest="disabled_syspreps",
106                       help="prevent SYSPREP operation from running on the "
107                       "input media", default=[], action="append",
108                       metavar="SYSPREP")
109
110     parser.add_option("--no-sysprep", dest="sysprep", default=True,
111                       help="don't perform any system preparation operation",
112                       action="store_false")
113
114     parser.add_option("--no-shrink", dest="shrink", default=True,
115                       help="don't shrink any partition", action="store_false")
116
117     parser.add_option("--public", dest="public", default=False,
118                       help="register image with cyclades as public",
119                       action="store_true")
120
121     parser.add_option("--tmpdir", dest="tmp", type="string", default=None,
122                       help="create large temporary image files under DIR",
123                       metavar="DIR")
124
125     options, args = parser.parse_args(input_args)
126
127     if len(args) != 1:
128         parser.error('Wrong number of arguments')
129
130     options.source = args[0]
131     if not os.path.exists(options.source):
132         raise FatalError("Input media `%s' is not accessible" % options.source)
133
134     if options.register and not options.upload:
135         raise FatalError("You also need to set -u when -r option is set")
136
137     if options.upload and options.token is None:
138         raise FatalError(
139             "Image uploading cannot be performed. "
140             "No authentication token is specified. Use -t to set a token")
141
142     if options.tmp is not None and not os.path.isdir(options.tmp):
143         raise FatalError("The directory `%s' specified with --tmpdir is not "
144                          "valid" % options.tmp)
145
146     meta = {}
147     for m in options.metadata:
148         try:
149             key, value = m.split('=', 1)
150         except ValueError:
151             raise FatalError("Metadata option: `%s' is not in KEY=VALUE "
152                              "format." % m)
153         meta[key] = value
154     options.metadata = meta
155
156     return options
157
158
159 def image_creator():
160     options = parse_options(sys.argv[1:])
161
162     if options.outfile is None and not options.upload and not \
163             options.print_sysprep:
164         raise FatalError("At least one of `-o', `-u' or `--print-sysprep' "
165                          "must be set")
166
167     if options.silent:
168         out = SilentOutput()
169     else:
170         out = OutputWthProgress(True) if sys.stderr.isatty() else \
171             SimpleOutput(False)
172
173     title = 'snf-image-creator %s' % version
174     out.output(title)
175     out.output('=' * len(title))
176
177     if os.geteuid() != 0:
178         raise FatalError("You must run %s as root"
179                          % os.path.basename(sys.argv[0]))
180
181     if not options.force and options.outfile is not None:
182         for extension in ('', '.meta', '.md5sum'):
183             filename = "%s%s" % (options.outfile, extension)
184             if os.path.exists(filename):
185                 raise FatalError("Output file %s exists "
186                                  "(use --force to overwrite it)" % filename)
187
188     # Check if the authentication token is valid. The earlier the better
189     if options.token is not None:
190         try:
191             account = Kamaki.get_account(options.token)
192             if account is None:
193                 raise FatalError("The authentication token you provided is not"
194                                  " valid!")
195         except ClientError as e:
196             raise FatalError("Astakos client: %d %s" % (e.status, e.message))
197
198     disk = Disk(options.source, out, options.tmp)
199
200     def signal_handler(signum, frame):
201         disk.cleanup()
202
203     signal.signal(signal.SIGINT, signal_handler)
204     signal.signal(signal.SIGTERM, signal_handler)
205     try:
206         snapshot = disk.snapshot()
207
208         dev = disk.get_device(snapshot)
209
210         # If no customization is to be applied, the image should be mounted ro
211         readonly = (not (options.sysprep or options.shrink) or
212                     options.print_sysprep)
213         dev.mount(readonly)
214
215         cls = os_cls(dev.distro, dev.ostype)
216         image_os = cls(dev.root, dev.g, out)
217         out.output()
218
219         for sysprep in options.disabled_syspreps:
220             image_os.disable_sysprep(image_os.get_sysprep_by_name(sysprep))
221
222         for sysprep in options.enabled_syspreps:
223             image_os.enable_sysprep(image_os.get_sysprep_by_name(sysprep))
224
225         if options.print_sysprep:
226             image_os.print_syspreps()
227             out.output()
228
229         if options.outfile is None and not options.upload:
230             return 0
231
232         if options.sysprep:
233             image_os.do_sysprep()
234
235         metadata = image_os.meta
236         dev.umount()
237
238         size = options.shrink and dev.shrink() or dev.size
239         metadata.update(dev.meta)
240
241         # Add command line metadata to the collected ones...
242         metadata.update(options.metadata)
243
244         md5 = MD5(out)
245         checksum = md5.compute(snapshot, size)
246
247         metastring = '\n'.join(
248             ['%s=%s' % (key, value) for (key, value) in metadata.items()])
249         metastring += '\n'
250
251         if options.outfile is not None:
252             dev.dump(options.outfile)
253
254             out.output('Dumping metadata file ...', False)
255             with open('%s.%s' % (options.outfile, 'meta'), 'w') as f:
256                 f.write(metastring)
257             out.success('done')
258
259             out.output('Dumping md5sum file ...', False)
260             with open('%s.%s' % (options.outfile, 'md5sum'), 'w') as f:
261                 f.write('%s %s\n' % (checksum,
262                                      os.path.basename(options.outfile)))
263             out.success('done')
264
265         # Destroy the device. We only need the snapshot from now on
266         disk.destroy_device(dev)
267
268         out.output()
269         try:
270             uploaded_obj = ""
271             if options.upload:
272                 out.output("Uploading image to pithos:")
273                 kamaki = Kamaki(account, out)
274                 with open(snapshot, 'rb') as f:
275                     uploaded_obj = kamaki.upload(
276                         f, size, options.upload,
277                         "(1/4)  Calculating block hashes",
278                         "(2/4)  Uploading missing blocks")
279
280                 out.output("(3/4)  Uploading metadata file ...", False)
281                 kamaki.upload(StringIO.StringIO(metastring),
282                               size=len(metastring),
283                               remote_path="%s.%s" % (options.upload, 'meta'))
284                 out.success('done')
285                 out.output("(4/4)  Uploading md5sum file ...", False)
286                 md5sumstr = '%s %s\n' % (checksum,
287                                          os.path.basename(options.upload))
288                 kamaki.upload(StringIO.StringIO(md5sumstr),
289                               size=len(md5sumstr),
290                               remote_path="%s.%s" % (options.upload, 'md5sum'))
291                 out.success('done')
292                 out.output()
293
294             if options.register:
295                 img_type = 'public' if options.public else 'private'
296                 out.output('Registering %s image with ~okeanos ...' % img_type,
297                            False)
298                 kamaki.register(options.register, uploaded_obj, metadata,
299                                 options.public)
300                 out.success('done')
301                 out.output()
302         except ClientError as e:
303             raise FatalError("Pithos client: %d %s" % (e.status, e.message))
304
305     finally:
306         out.output('cleaning up ...')
307         disk.cleanup()
308
309     out.success("snf-image-creator exited without errors")
310
311     return 0
312
313
314 def main():
315     try:
316         ret = image_creator()
317         sys.exit(ret)
318     except FatalError as e:
319         colored = sys.stderr.isatty()
320         SimpleOutput(colored).error(e)
321         sys.exit(1)
322
323 if __name__ == '__main__':
324     main()
325
326 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :