Statistics
| Branch: | Tag: | Revision:

root / image_creator / main.py @ 13675429

History | View | Annotate | Download (6.5 kB)

1 ae48a082 Nikos Skalkotos
#!/usr/bin/env python
2 ae48a082 Nikos Skalkotos
3 d57775d4 Nikos Skalkotos
# Copyright 2011 GRNET S.A. All rights reserved.
4 d57775d4 Nikos Skalkotos
#
5 d57775d4 Nikos Skalkotos
# Redistribution and use in source and binary forms, with or
6 d57775d4 Nikos Skalkotos
# without modification, are permitted provided that the following
7 d57775d4 Nikos Skalkotos
# conditions are met:
8 d57775d4 Nikos Skalkotos
#
9 d57775d4 Nikos Skalkotos
#   1. Redistributions of source code must retain the above
10 d57775d4 Nikos Skalkotos
#      copyright notice, this list of conditions and the following
11 d57775d4 Nikos Skalkotos
#      disclaimer.
12 d57775d4 Nikos Skalkotos
#
13 d57775d4 Nikos Skalkotos
#   2. Redistributions in binary form must reproduce the above
14 d57775d4 Nikos Skalkotos
#      copyright notice, this list of conditions and the following
15 d57775d4 Nikos Skalkotos
#      disclaimer in the documentation and/or other materials
16 d57775d4 Nikos Skalkotos
#      provided with the distribution.
17 d57775d4 Nikos Skalkotos
#
18 d57775d4 Nikos Skalkotos
# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
19 d57775d4 Nikos Skalkotos
# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
20 d57775d4 Nikos Skalkotos
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
21 d57775d4 Nikos Skalkotos
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
22 d57775d4 Nikos Skalkotos
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23 d57775d4 Nikos Skalkotos
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24 d57775d4 Nikos Skalkotos
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
25 d57775d4 Nikos Skalkotos
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
26 d57775d4 Nikos Skalkotos
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
27 d57775d4 Nikos Skalkotos
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
28 d57775d4 Nikos Skalkotos
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 d57775d4 Nikos Skalkotos
# POSSIBILITY OF SUCH DAMAGE.
30 d57775d4 Nikos Skalkotos
#
31 d57775d4 Nikos Skalkotos
# The views and conclusions contained in the software and
32 d57775d4 Nikos Skalkotos
# documentation are those of the authors and should not be
33 d57775d4 Nikos Skalkotos
# interpreted as representing official policies, either expressed
34 d57775d4 Nikos Skalkotos
# or implied, of GRNET S.A.
35 d57775d4 Nikos Skalkotos
36 aa2062ba Nikos Skalkotos
from image_creator import get_os_class
37 c408053f Nikos Skalkotos
from image_creator import __version__ as version
38 0ae01e26 Nikos Skalkotos
from image_creator import FatalError
39 d57775d4 Nikos Skalkotos
from image_creator.disk import Disk
40 22a6d232 Nikos Skalkotos
from image_creator.util import get_command, error, progress_generator, success
41 22a6d232 Nikos Skalkotos
from clint.textui import puts, indent
42 22a6d232 Nikos Skalkotos
from sendfile import sendfile
43 3ccb2618 Nikos Skalkotos
44 d57775d4 Nikos Skalkotos
import sys
45 d57775d4 Nikos Skalkotos
import os
46 c408053f Nikos Skalkotos
import optparse
47 3ccb2618 Nikos Skalkotos
48 3ccb2618 Nikos Skalkotos
dd = get_command('dd')
49 c408053f Nikos Skalkotos
50 c408053f Nikos Skalkotos
51 0ae01e26 Nikos Skalkotos
def check_writable_dir(option, opt_str, value, parser):
52 0ae01e26 Nikos Skalkotos
    dirname = os.path.dirname(value)
53 0ae01e26 Nikos Skalkotos
    name = os.path.basename(value)
54 0ae01e26 Nikos Skalkotos
    if dirname and not os.path.isdir(dirname):
55 0ae01e26 Nikos Skalkotos
        parser.error("`%s' is not an existing directory" % dirname)
56 c408053f Nikos Skalkotos
57 0ae01e26 Nikos Skalkotos
    if not name:
58 0ae01e26 Nikos Skalkotos
        parser.error("`%s' is not a valid file name" % dirname)
59 c408053f Nikos Skalkotos
60 c408053f Nikos Skalkotos
    setattr(parser.values, option.dest, value)
61 c408053f Nikos Skalkotos
62 c408053f Nikos Skalkotos
63 c408053f Nikos Skalkotos
def parse_options(input_args):
64 0ae01e26 Nikos Skalkotos
    usage = "Usage: %prog [options] <input_media>"
65 c408053f Nikos Skalkotos
    parser = optparse.OptionParser(version=version, usage=usage)
66 c408053f Nikos Skalkotos
67 c408053f Nikos Skalkotos
    parser.add_option("-f", "--force", dest="force", default=False,
68 c408053f Nikos Skalkotos
        action="store_true", help="Overwrite output files if they exist")
69 c408053f Nikos Skalkotos
70 1a3f1298 Nikos Skalkotos
    parser.add_option("--no-cleanup", dest="cleanup", default=True,
71 1a3f1298 Nikos Skalkotos
        help="Don't cleanup sensitive data before extracting the image",
72 1a3f1298 Nikos Skalkotos
        action="store_false")
73 1a3f1298 Nikos Skalkotos
74 9cbb5794 Nikos Skalkotos
    parser.add_option("--no-sysprep", dest="sysprep", default=True,
75 9cbb5794 Nikos Skalkotos
        help="Don't perform system preperation before extracting the image",
76 9cbb5794 Nikos Skalkotos
        action="store_false")
77 9cbb5794 Nikos Skalkotos
78 9cbb5794 Nikos Skalkotos
    parser.add_option("--no-shrink", dest="shrink", default=True,
79 9cbb5794 Nikos Skalkotos
        help="Don't shrink any partition before extracting the image",
80 9cbb5794 Nikos Skalkotos
        action="store_false")
81 9cbb5794 Nikos Skalkotos
82 0ae01e26 Nikos Skalkotos
    parser.add_option("-o", "--outfile", type="string", dest="outfile",
83 0ae01e26 Nikos Skalkotos
        default=None, action="callback", callback=check_writable_dir,
84 0ae01e26 Nikos Skalkotos
        help="Output image file",
85 0ae01e26 Nikos Skalkotos
        metavar="FILE")
86 0ae01e26 Nikos Skalkotos
87 1a3f1298 Nikos Skalkotos
    parser.add_option("-u", "--upload", dest="upload", default=False,
88 1a3f1298 Nikos Skalkotos
        help="Upload image to a pithos repository using kamaki",
89 1a3f1298 Nikos Skalkotos
        action="store_true")
90 1a3f1298 Nikos Skalkotos
91 1a3f1298 Nikos Skalkotos
    parser.add_option("-r", "--register", dest="register", default=False,
92 1a3f1298 Nikos Skalkotos
        help="Register image to okeanos using kamaki", action="store_true")
93 1a3f1298 Nikos Skalkotos
94 c408053f Nikos Skalkotos
    options, args = parser.parse_args(input_args)
95 c408053f Nikos Skalkotos
96 0ae01e26 Nikos Skalkotos
    if len(args) != 1:
97 0ae01e26 Nikos Skalkotos
        parser.error('Wrong number of arguments')
98 c408053f Nikos Skalkotos
    options.source = args[0]
99 c408053f Nikos Skalkotos
    if not os.path.exists(options.source):
100 0ae01e26 Nikos Skalkotos
        parser.error('input media is not accessible')
101 c408053f Nikos Skalkotos
102 1a3f1298 Nikos Skalkotos
    if options.register:
103 1a3f1298 Nikos Skalkotos
        options.upload = True
104 1a3f1298 Nikos Skalkotos
105 0ae01e26 Nikos Skalkotos
    if options.outfile is None and not options.upload:
106 0ae01e26 Nikos Skalkotos
        parser.error('either outfile (-o) or upload (-u) must be set.')
107 0ae01e26 Nikos Skalkotos
108 c408053f Nikos Skalkotos
    return options
109 d57775d4 Nikos Skalkotos
110 8c574358 Nikos Skalkotos
111 22a6d232 Nikos Skalkotos
def extract_image(device, outfile, size):
112 22a6d232 Nikos Skalkotos
    blocksize = 4194304  # 4MB
113 22a6d232 Nikos Skalkotos
    progress_size = (size + 1048575) // 1048576  # in MB
114 22a6d232 Nikos Skalkotos
    progressbar = progress_generator("Dumping image file: ",
115 22a6d232 Nikos Skalkotos
                                                    progress_size)
116 22a6d232 Nikos Skalkotos
    source = open(device, "r")
117 22a6d232 Nikos Skalkotos
    try:
118 22a6d232 Nikos Skalkotos
        dest = open(outfile, "w")
119 22a6d232 Nikos Skalkotos
        try:
120 22a6d232 Nikos Skalkotos
            left = size
121 22a6d232 Nikos Skalkotos
            offset = 0
122 22a6d232 Nikos Skalkotos
            progressbar.next()
123 22a6d232 Nikos Skalkotos
            while left > 0:
124 22a6d232 Nikos Skalkotos
                length = min(left, blocksize)
125 22a6d232 Nikos Skalkotos
                sent = sendfile(dest.fileno(), source.fileno(), offset, length)
126 22a6d232 Nikos Skalkotos
                offset += sent
127 22a6d232 Nikos Skalkotos
                left -= sent
128 22a6d232 Nikos Skalkotos
                for i in range(4):
129 22a6d232 Nikos Skalkotos
                    progressbar.next()
130 22a6d232 Nikos Skalkotos
        finally:
131 22a6d232 Nikos Skalkotos
            dest.close()
132 22a6d232 Nikos Skalkotos
    finally:
133 22a6d232 Nikos Skalkotos
        source.close()
134 22a6d232 Nikos Skalkotos
135 22a6d232 Nikos Skalkotos
    success('Image file %s was successfully created' % outfile)
136 22a6d232 Nikos Skalkotos
137 d57775d4 Nikos Skalkotos
138 22a6d232 Nikos Skalkotos
def image_creator():
139 22a6d232 Nikos Skalkotos
    puts('snf-image-creator %s\n' % version)
140 c408053f Nikos Skalkotos
    options = parse_options(sys.argv[1:])
141 c408053f Nikos Skalkotos
142 c408053f Nikos Skalkotos
    if os.geteuid() != 0:
143 c408053f Nikos Skalkotos
        raise FatalError("You must run %s as root" \
144 c408053f Nikos Skalkotos
                        % os.path.basename(sys.argv[0]))
145 c408053f Nikos Skalkotos
146 c408053f Nikos Skalkotos
    if not options.force:
147 0ae01e26 Nikos Skalkotos
        for extension in ('', '.meta'):
148 0ae01e26 Nikos Skalkotos
            filename = "%s%s" % (options.outfile, extension)
149 c408053f Nikos Skalkotos
            if os.path.exists(filename):
150 c408053f Nikos Skalkotos
                raise FatalError("Output file %s exists "
151 c408053f Nikos Skalkotos
                    "(use --force to overwrite it)." % filename)
152 c408053f Nikos Skalkotos
153 c408053f Nikos Skalkotos
    disk = Disk(options.source)
154 d57775d4 Nikos Skalkotos
    try:
155 d57775d4 Nikos Skalkotos
        dev = disk.get_device()
156 0d5a999d Nikos Skalkotos
        dev.mount()
157 22a6d232 Nikos Skalkotos
158 aa2062ba Nikos Skalkotos
        osclass = get_os_class(dev.distro, dev.ostype)
159 aa2062ba Nikos Skalkotos
        image_os = osclass(dev.root, dev.g)
160 aa2062ba Nikos Skalkotos
        metadata = image_os.get_metadata()
161 9cbb5794 Nikos Skalkotos
162 3f70f242 Nikos Skalkotos
        puts()
163 3f70f242 Nikos Skalkotos
164 9cbb5794 Nikos Skalkotos
        if options.sysprep:
165 9cbb5794 Nikos Skalkotos
            image_os.sysprep()
166 ae48a082 Nikos Skalkotos
167 1a3f1298 Nikos Skalkotos
        if options.cleanup:
168 1a3f1298 Nikos Skalkotos
            image_os.data_cleanup()
169 1a3f1298 Nikos Skalkotos
170 8c574358 Nikos Skalkotos
        dev.umount()
171 1a3f1298 Nikos Skalkotos
172 c408053f Nikos Skalkotos
        size = options.shrink and dev.shrink() or dev.size()
173 c408053f Nikos Skalkotos
        metadata['size'] = str(size // 2 ** 20)
174 ae48a082 Nikos Skalkotos
175 0ae01e26 Nikos Skalkotos
        if options.outfile is not None:
176 0ae01e26 Nikos Skalkotos
            f = open('%s.%s' % (options.outfile, 'meta'), 'w')
177 0ae01e26 Nikos Skalkotos
            try:
178 0ae01e26 Nikos Skalkotos
                for key in metadata.keys():
179 0ae01e26 Nikos Skalkotos
                    f.write("%s=%s\n" % (key, metadata[key]))
180 0ae01e26 Nikos Skalkotos
            finally:
181 0ae01e26 Nikos Skalkotos
                f.close()
182 0ae01e26 Nikos Skalkotos
183 22a6d232 Nikos Skalkotos
            extract_image(dev.device, options.outfile, size)
184 d57775d4 Nikos Skalkotos
    finally:
185 22a6d232 Nikos Skalkotos
        puts('cleaning up...')
186 d57775d4 Nikos Skalkotos
        disk.cleanup()
187 d57775d4 Nikos Skalkotos
188 c408053f Nikos Skalkotos
    return 0
189 c408053f Nikos Skalkotos
190 ae48a082 Nikos Skalkotos
191 0ae01e26 Nikos Skalkotos
def main():
192 c408053f Nikos Skalkotos
    try:
193 0ae01e26 Nikos Skalkotos
        ret = image_creator()
194 c408053f Nikos Skalkotos
        sys.exit(ret)
195 c408053f Nikos Skalkotos
    except FatalError as e:
196 22a6d232 Nikos Skalkotos
        error(e)
197 c408053f Nikos Skalkotos
        sys.exit(1)
198 d57775d4 Nikos Skalkotos
199 0ae01e26 Nikos Skalkotos
200 0ae01e26 Nikos Skalkotos
if __name__ == '__main__':
201 0ae01e26 Nikos Skalkotos
    main()
202 ae48a082 Nikos Skalkotos
203 d57775d4 Nikos Skalkotos
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :