Statistics
| Branch: | Tag: | Revision:

root / image_creator / main.py @ b1395967

History | View | Annotate | Download (9.5 kB)

1
#!/usr/bin/env python
2

    
3
# Copyright 2011 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 get_os_class
37
from image_creator import __version__ as version
38
from image_creator import util
39
from image_creator.disk import Disk
40
from image_creator.util import get_command, error, success, output, \
41
                                                    FatalError, progress, md5
42
from image_creator.kamaki_wrapper import Kamaki
43
import sys
44
import os
45
import optparse
46
import StringIO
47

    
48
dd = get_command('dd')
49

    
50

    
51
def check_writable_dir(option, opt_str, value, parser):
52
    dirname = os.path.dirname(value)
53
    name = os.path.basename(value)
54
    if dirname and not os.path.isdir(dirname):
55
        raise FatalError("`%s' is not an existing directory" % dirname)
56

    
57
    if not name:
58
        raise FatalError("`%s' is not a valid file name" % dirname)
59

    
60
    setattr(parser.values, option.dest, value)
61

    
62

    
63
def parse_options(input_args):
64
    usage = "Usage: %prog [options] <input_media>"
65
    parser = optparse.OptionParser(version=version, usage=usage)
66

    
67
    account = os.environ["OKEANOS_USER"] if "OKEANOS_USER" in os.environ \
68
        else None
69
    token = os.environ["OKEANOS_TOKEN"] if "OKEANOS_TOKEN" in os.environ \
70
        else None
71

    
72
    parser.add_option("-o", "--outfile", type="string", dest="outfile",
73
        default=None, action="callback", callback=check_writable_dir,
74
        help="dump image to FILE", metavar="FILE")
75

    
76
    parser.add_option("-f", "--force", dest="force", default=False,
77
        action="store_true", help="overwrite output files if they exist")
78

    
79
    parser.add_option("-s", "--silent", dest="silent", default=False,
80
        help="silent mode, only output errors", action="store_true")
81

    
82
    parser.add_option("-u", "--upload", dest="upload", type="string",
83
        default=False, help="upload the image to pithos with name FILENAME",
84
        metavar="FILENAME")
85

    
86
    parser.add_option("-r", "--register", dest="register", type="string",
87
        default=False, help="register the image to ~okeanos as IMAGENAME",
88
        metavar="IMAGENAME")
89

    
90
    parser.add_option("-a", "--account", dest="account", type="string",
91
        default=account,
92
        help="Use this ACCOUNT when uploading/registring images [Default: %s]"\
93
        % account)
94

    
95
    parser.add_option("-t", "--token", dest="token", type="string",
96
        default=token,
97
        help="Use this token when uploading/registring images [Default: %s]"\
98
        % token)
99

    
100
    parser.add_option("--print-sysprep", dest="print_sysprep", default=False,
101
        help="print the enabled and disabled system preparation operations "
102
        "for this input media", action="store_true")
103

    
104
    parser.add_option("--enable-sysprep", dest="enabled_syspreps", default=[],
105
        help="run SYSPREP operation on the input media",
106
        action="append", metavar="SYSPREP")
107

    
108
    parser.add_option("--disable-sysprep", dest="disabled_syspreps",
109
        help="prevent SYSPREP operation from running on the input media",
110
        default=[], action="append", metavar="SYSPREP")
111

    
112
    parser.add_option("--no-sysprep", dest="sysprep", default=True,
113
        help="don't perform system preperation", action="store_false")
114

    
115
    parser.add_option("--no-shrink", dest="shrink", default=True,
116
        help="don't shrink any partition", action="store_false")
117

    
118
    options, args = parser.parse_args(input_args)
119

    
120
    if len(args) != 1:
121
        parser.error('Wrong number of arguments')
122
    options.source = args[0]
123
    if not os.path.exists(options.source):
124
        raise FatalError("Input media `%s' is not accessible" % options.source)
125

    
126
    if options.register and options.upload == False:
127
        raise FatalError("You also need to set -u when -r option is set")
128

    
129
    if options.upload and options.account is None:
130
        raise FatalError("Image uploading cannot be performed. No ~okeanos "
131
        "account name is specified. Use -a to set an account name.")
132

    
133
    if options.upload and options.token is None:
134
        raise FatalError("Image uploading cannot be performed. No ~okeanos "
135
        "token is specified. User -t to set a token.")
136

    
137
    return options
138

    
139

    
140
def image_creator():
141
    options = parse_options(sys.argv[1:])
142

    
143
    if options.silent:
144
        util.silent = True
145

    
146
    if options.outfile is None and not options.upload \
147
                                            and not options.print_sysprep:
148
        raise FatalError("At least one of `-o', `-u' or" \
149
                            "`--print-sysprep' must be set")
150

    
151
    output('snf-image-creator %s\n' % version)
152

    
153
    if os.geteuid() != 0:
154
        raise FatalError("You must run %s as root" \
155
                        % os.path.basename(sys.argv[0]))
156

    
157
    if not options.force and options.outfile is not None:
158
        for extension in ('', '.meta', '.md5sum'):
159
            filename = "%s%s" % (options.outfile, extension)
160
            if os.path.exists(filename):
161
                raise FatalError("Output file %s exists "
162
                    "(use --force to overwrite it)." % filename)
163

    
164
    disk = Disk(options.source)
165
    try:
166
        snapshot = disk.snapshot()
167

    
168
        dev = disk.get_device(snapshot)
169
        dev.mount()
170

    
171
        osclass = get_os_class(dev.distro, dev.ostype)
172
        image_os = osclass(dev.root, dev.g)
173
        metadata = image_os.get_metadata()
174

    
175
        output()
176

    
177
        for sysprep in options.disabled_syspreps:
178
            image_os.disable_sysprep(sysprep)
179

    
180
        for sysprep in options.enabled_syspreps:
181
            image_os.enable_sysprep(sysprep)
182

    
183
        if options.print_sysprep:
184
            image_os.print_syspreps()
185
            output()
186

    
187
        if options.outfile is None and not options.upload:
188
            return 0
189

    
190
        if options.sysprep:
191
            image_os.do_sysprep()
192

    
193
        dev.umount()
194

    
195
        size = options.shrink and dev.shrink() or dev.size()
196
        metadata['SIZE'] = str(size // 2 ** 20)
197

    
198
        checksum = md5(snapshot, size)
199

    
200
        metastring = "\n".join(
201
            ['%s=%s' % (key, value) for (key, value) in metadata.items()])
202

    
203
        if options.outfile is not None:
204
            dev.dump(options.outfile)
205

    
206
            output('Dumping metadata file...', False)
207
            with open('%s.%s' % (options.outfile, 'meta'), 'w') as f:
208
                    f.write(metastring)
209
            success('done')
210

    
211
            output('Dumping md5sum file...', False)
212
            with open('%s.%s' % (options.outfile, 'md5sum'), 'w') as f:
213
                f.write('%s %s'% (checksum, os.path.basename(options.outfile)))
214
            success('done')
215

    
216
        # Destroy the device. We only need the snapshot from now on
217
        disk.destroy_device(dev)
218

    
219
        output()
220

    
221
        uploaded_obj = ""
222
        if options.upload:
223
            output("Uploading image to pithos:")
224
            kamaki = Kamaki(options.account, options.token)
225
            with open(snapshot) as f:
226
                uploaded_obj = kamaki.upload(f, size, options.upload,
227
                                "(1/4)  Calculating block hashes",
228
                                "(2/4)  Uploading missing blocks")
229

    
230
            output("(3/4)  Uploading metadata file...", False)
231
            kamaki.upload(StringIO.StringIO(metastring), size=len(metastring),
232
                                remote_path="%s.%s" % (options.upload, 'meta'))
233
            success('done')
234
            output("(4/4)  Uploading md5sum file...", False)
235
            md5sumstr = '%s %s' % (checksum, os.path.basename(options.upload))
236
            kamaki.upload(StringIO.StringIO(md5sumstr), size=len(md5sumstr),
237
                            remote_path="%s.%s" % (options.upload, 'md5sum'))
238
            success('done')
239
            output()
240

    
241
        if options.register:
242
            output('Registing image to ~okeanos...')
243
            kamaki.register(options.register, uploaded_obj, metadata)
244
            output('done')
245
            output()
246

    
247
    finally:
248
        output('cleaning up...')
249
        disk.cleanup()
250

    
251
    success("snf-image-creator exited without errors")
252

    
253
    return 0
254

    
255

    
256
def main():
257
    try:
258
        ret = image_creator()
259
        sys.exit(ret)
260
    except FatalError as e:
261
        error(e)
262
        sys.exit(1)
263

    
264

    
265
if __name__ == '__main__':
266
    main()
267

    
268
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :