Statistics
| Branch: | Tag: | Revision:

root / image_creator / main.py @ 13675429

History | View | Annotate | Download (6.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 FatalError
39
from image_creator.disk import Disk
40
from image_creator.util import get_command, error, progress_generator, success
41
from clint.textui import puts, indent
42
from sendfile import sendfile
43

    
44
import sys
45
import os
46
import optparse
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
        parser.error("`%s' is not an existing directory" % dirname)
56

    
57
    if not name:
58
        parser.error("`%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
    parser.add_option("-f", "--force", dest="force", default=False,
68
        action="store_true", help="Overwrite output files if they exist")
69

    
70
    parser.add_option("--no-cleanup", dest="cleanup", default=True,
71
        help="Don't cleanup sensitive data before extracting the image",
72
        action="store_false")
73

    
74
    parser.add_option("--no-sysprep", dest="sysprep", default=True,
75
        help="Don't perform system preperation before extracting the image",
76
        action="store_false")
77

    
78
    parser.add_option("--no-shrink", dest="shrink", default=True,
79
        help="Don't shrink any partition before extracting the image",
80
        action="store_false")
81

    
82
    parser.add_option("-o", "--outfile", type="string", dest="outfile",
83
        default=None, action="callback", callback=check_writable_dir,
84
        help="Output image file",
85
        metavar="FILE")
86

    
87
    parser.add_option("-u", "--upload", dest="upload", default=False,
88
        help="Upload image to a pithos repository using kamaki",
89
        action="store_true")
90

    
91
    parser.add_option("-r", "--register", dest="register", default=False,
92
        help="Register image to okeanos using kamaki", action="store_true")
93

    
94
    options, args = parser.parse_args(input_args)
95

    
96
    if len(args) != 1:
97
        parser.error('Wrong number of arguments')
98
    options.source = args[0]
99
    if not os.path.exists(options.source):
100
        parser.error('input media is not accessible')
101

    
102
    if options.register:
103
        options.upload = True
104

    
105
    if options.outfile is None and not options.upload:
106
        parser.error('either outfile (-o) or upload (-u) must be set.')
107

    
108
    return options
109

    
110

    
111
def extract_image(device, outfile, size):
112
    blocksize = 4194304  # 4MB
113
    progress_size = (size + 1048575) // 1048576  # in MB
114
    progressbar = progress_generator("Dumping image file: ",
115
                                                    progress_size)
116
    source = open(device, "r")
117
    try:
118
        dest = open(outfile, "w")
119
        try:
120
            left = size
121
            offset = 0
122
            progressbar.next()
123
            while left > 0:
124
                length = min(left, blocksize)
125
                sent = sendfile(dest.fileno(), source.fileno(), offset, length)
126
                offset += sent
127
                left -= sent
128
                for i in range(4):
129
                    progressbar.next()
130
        finally:
131
            dest.close()
132
    finally:
133
        source.close()
134

    
135
    success('Image file %s was successfully created' % outfile)
136

    
137

    
138
def image_creator():
139
    puts('snf-image-creator %s\n' % version)
140
    options = parse_options(sys.argv[1:])
141

    
142
    if os.geteuid() != 0:
143
        raise FatalError("You must run %s as root" \
144
                        % os.path.basename(sys.argv[0]))
145

    
146
    if not options.force:
147
        for extension in ('', '.meta'):
148
            filename = "%s%s" % (options.outfile, extension)
149
            if os.path.exists(filename):
150
                raise FatalError("Output file %s exists "
151
                    "(use --force to overwrite it)." % filename)
152

    
153
    disk = Disk(options.source)
154
    try:
155
        dev = disk.get_device()
156
        dev.mount()
157

    
158
        osclass = get_os_class(dev.distro, dev.ostype)
159
        image_os = osclass(dev.root, dev.g)
160
        metadata = image_os.get_metadata()
161

    
162
        puts()
163

    
164
        if options.sysprep:
165
            image_os.sysprep()
166

    
167
        if options.cleanup:
168
            image_os.data_cleanup()
169

    
170
        dev.umount()
171

    
172
        size = options.shrink and dev.shrink() or dev.size()
173
        metadata['size'] = str(size // 2 ** 20)
174

    
175
        if options.outfile is not None:
176
            f = open('%s.%s' % (options.outfile, 'meta'), 'w')
177
            try:
178
                for key in metadata.keys():
179
                    f.write("%s=%s\n" % (key, metadata[key]))
180
            finally:
181
                f.close()
182

    
183
            extract_image(dev.device, options.outfile, size)
184
    finally:
185
        puts('cleaning up...')
186
        disk.cleanup()
187

    
188
    return 0
189

    
190

    
191
def main():
192
    try:
193
        ret = image_creator()
194
        sys.exit(ret)
195
    except FatalError as e:
196
        error(e)
197
        sys.exit(1)
198

    
199

    
200
if __name__ == '__main__':
201
    main()
202

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