Revision 22a6d232

b/image_creator/disk.py
31 31
# interpreted as representing official policies, either expressed
32 32
# or implied, of GRNET S.A.
33 33

  
34
from image_creator.util import get_command
34
from image_creator.util import get_command, warn, progress_generator
35 35
from image_creator import FatalError
36
from clint.textui import progress
36
from clint.textui import indent, puts, colored
37 37

  
38 38
import stat
39 39
import os
......
99 99
        This instance is a snapshot of the original source media of
100 100
        the Disk instance.
101 101
        """
102
        sourcedev = self.source
103
        mode = os.stat(self.source).st_mode
104
        if stat.S_ISDIR(mode):
105
            return self._losetup(self._dir_to_disk())
106
        elif stat.S_ISREG(mode):
107
            sourcedev = self._losetup(self.source)
108
        elif not stat.S_ISBLK(mode):
109
            raise ValueError("Value for self.source is invalid")
102

  
103
        puts("Examining source media `%s'" % self.source)
104
        with indent(4):
105
            sourcedev = self.source
106
            mode = os.stat(self.source).st_mode
107
            if stat.S_ISDIR(mode):
108
                puts(colored.green('Looks like a directory'))
109
                return self._losetup(self._dir_to_disk())
110
            elif stat.S_ISREG(mode):
111
                puts(colored.green('Looks like an image file'))
112
                sourcedev = self._losetup(self.source)
113
            elif not stat.S_ISBLK(mode):
114
                raise ValueError("Invalid media source. Only block devices, "
115
                                "regular files and directories are supported.")
116
            else:
117
                puts(colored.green('Looks like a block device'))
118
            #puts()
110 119

  
111 120
        # Take a snapshot and return it to the user
112
        size = blockdev('--getsize', sourcedev)
113
        cowfd, cow = tempfile.mkstemp()
114
        os.close(cowfd)
115
        self._add_cleanup(os.unlink, cow)
116
        # Create 1G cow sparse file
117
        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
118
        cowdev = self._losetup(cow)
119

  
120
        snapshot = uuid.uuid4().hex
121
        tablefd, table = tempfile.mkstemp()
122
        try:
123
            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
124
                                        (int(size), sourcedev, cowdev))
125
            dmsetup('create', snapshot, table)
126
            self._add_cleanup(dmsetup, 'remove', snapshot)
127
        finally:
128
            os.unlink(table)
121
        puts("Snapshotting media source")
122
        with indent(4):
123
            size = blockdev('--getsize', sourcedev)
124
            cowfd, cow = tempfile.mkstemp()
125
            os.close(cowfd)
126
            self._add_cleanup(os.unlink, cow)
127
            # Create 1G cow sparse file
128
            dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
129
                                            'seek=%d' % (1024 * 1024))
130
            cowdev = self._losetup(cow)
131
    
132
            snapshot = uuid.uuid4().hex
133
            tablefd, table = tempfile.mkstemp()
134
            try:
135
                os.write(tablefd, "0 %d snapshot %s %s n 8" % \
136
                                            (int(size), sourcedev, cowdev))
137
                dmsetup('create', snapshot, table)
138
                self._add_cleanup(dmsetup, 'remove', snapshot)
139
                # Sometimes dmsetup remove fails with Device or resource busy,
140
                # although everything is cleaned up and the snapshot is not
141
                # used by anyone. Add a 2 seconds delay to be on the safe side.
142
                self._add_cleanup(time.sleep, 2)
143

  
144
            finally:
145
                os.unlink(table)
146
            puts(colored.green('Done'))
147
        # puts()
129 148
        new_device = DiskDevice("/dev/mapper/%s" % snapshot)
130 149
        self._devices.append(new_device)
131 150
        new_device.enable()
......
139 158
        device.destroy()
140 159

  
141 160

  
142
def progress_generator(label=''):
143
    position = 0
144
    for i in progress.bar(range(100), label):
145
        if i < position:
146
            continue
147
        position = yield
148
    yield  # suppress the StopIteration exception
149

  
150

  
151 161
class DiskDevice(object):
152 162
    """This class represents a block device hosting an Operating System
153 163
    as created by the device-mapper.
......
170 180

  
171 181
    def enable(self):
172 182
        """Enable a newly created DiskDevice"""
173

  
174
        self.progressbar = progress_generator("VM lauch: ")
175
        self.progressbar.next()
176
        eh = self.g.set_event_callback(self.progress_callback,
183
        self.progressbar = progress_generator("Launching helper VM: ")
184
        with indent(4):
185
            self.progressbar.next()
186
            eh = self.g.set_event_callback(self.progress_callback,
177 187
                                                        guestfs.EVENT_PROGRESS)
178
        self.g.launch()
179
        self.guestfs_enabled = True
180
        self.g.delete_event_callback(eh)
181
        if self.progressbar is not None:
182
            self.progressbar.send(100)
183
            self.progressbar = None
184

  
185
        roots = self.g.inspect_os()
186
        if len(roots) == 0:
187
            raise FatalError("No operating system found")
188
        if len(roots) > 1:
189
            raise FatalError("Multiple operating systems found")
190

  
191
        self.root = roots[0]
192
        self.ostype = self.g.inspect_get_type(self.root)
193
        self.distro = self.g.inspect_get_distro(self.root)
188
            self.g.launch()
189
            self.guestfs_enabled = True
190
            self.g.delete_event_callback(eh)
191
            if self.progressbar is not None:
192
                self.progressbar.send(100)
193
                self.progressbar = None
194
            puts(colored.green('Done'))
195

  
196
        puts('Inspecting Operating System')
197
        with indent(4):
198
            roots = self.g.inspect_os()
199
            if len(roots) == 0:
200
                raise FatalError("No operating system found")
201
            if len(roots) > 1:
202
                raise FatalError("Multiple operating systems found."
203
                                "We only support images with one filesystem.")
204
            self.root = roots[0]
205
            self.ostype = self.g.inspect_get_type(self.root)
206
            self.distro = self.g.inspect_get_distro(self.root)
207
            puts(colored.green('Found a %s system' % self.distro))
208
        puts()
194 209

  
195 210
    def destroy(self):
196 211
        """Destroy this DiskDevice instance."""
......
242 257
        disk and then updating the partition table. The new disk size
243 258
        (in bytes) is returned.
244 259
        """
260
        puts("Shrinking image (this may take a while)")
261
        
245 262
        dev = self.g.part_to_dev(self.root)
246 263
        parttype = self.g.part_get_parttype(dev)
247 264
        if parttype != 'msdos':
......
257 274
        part_dev = "%s%d" % (dev, last_partition['part_num'])
258 275
        fs_type = self.g.vfs_type(part_dev)
259 276
        if not re.match("ext[234]", fs_type):
260
            print "Warning: Don't know how to resize %s partitions." % vfs_type
277
            warn("Don't know how to resize %s partitions." % vfs_type)
261 278
            return
262 279

  
263
        self.g.e2fsck_f(part_dev)
264
        self.g.resize2fs_M(part_dev)
265
        output = self.g.tune2fs_l(part_dev)
266
        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
267
        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
280
        with indent(4):
281
            self.g.e2fsck_f(part_dev)
282
            self.g.resize2fs_M(part_dev)
283

  
284
            output = self.g.tune2fs_l(part_dev)
285
            block_size = int(
286
                filter(lambda x: x[0] == 'Block size', output)[0][1])
287
            block_cnt = int(
288
                filter(lambda x: x[0] == 'Block count', output)[0][1])
268 289

  
269
        sector_size = self.g.blockdev_getss(dev)
290
            sector_size = self.g.blockdev_getss(dev)
270 291

  
271
        start = last_partition['part_start'] / sector_size
272
        end = start + (block_size * block_cnt) / sector_size - 1
292
            start = last_partition['part_start'] / sector_size
293
            end = start + (block_size * block_cnt) / sector_size - 1
273 294

  
274
        self.g.part_del(dev, last_partition['part_num'])
275
        self.g.part_add(dev, 'p', start, end)
295
            self.g.part_del(dev, last_partition['part_num'])
296
            self.g.part_add(dev, 'p', start, end)
276 297

  
277
        return (end + 1) * sector_size
298
            new_size = (end + 1) * sector_size
299
            puts("  New image size is %dMB\n" % (new_size // 2 ** 20))
300
        return new_size
278 301

  
279 302
    def size(self):
280 303
        """Returns the "payload" size of the device.
b/image_creator/main.py
37 37
from image_creator import __version__ as version
38 38
from image_creator import FatalError
39 39
from image_creator.disk import Disk
40
from image_creator.util import get_command
40
from image_creator.util import get_command, error, progress_generator, success
41
from clint.textui import puts, indent
42
from sendfile import sendfile
41 43

  
42 44
import sys
43 45
import os
......
106 108
    return options
107 109

  
108 110

  
109
def image_creator():
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

  
110 137

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

  
113 142
    if os.geteuid() != 0:
......
125 154
    try:
126 155
        dev = disk.get_device()
127 156
        dev.mount()
157

  
128 158
        osclass = get_os_class(dev.distro, dev.ostype)
129 159
        image_os = osclass(dev.root, dev.g)
130 160
        metadata = image_os.get_metadata()
......
140 170
        size = options.shrink and dev.shrink() or dev.size()
141 171
        metadata['size'] = str(size // 2 ** 20)
142 172

  
143
        outfile = ""
144 173
        if options.outfile is not None:
145
            outfile = options.outfile
146 174
            f = open('%s.%s' % (options.outfile, 'meta'), 'w')
147 175
            try:
148 176
                for key in metadata.keys():
149 177
                    f.write("%s=%s\n" % (key, metadata[key]))
150 178
            finally:
151 179
                f.close()
152
        else:
153
            outfd, outfile = tmpfile.mkstemp()
154
            os.close(outfd)
155 180

  
156
        dd('if=%s' % dev.device,
157
            'of=%s' % outfile,
158
            'bs=4M', 'count=%d' % ((size + 1) // 2 ** 22))
181
            extract_image(dev.device, options.outfile, size)
159 182

  
160 183
    finally:
184
        puts('cleaning up...')
161 185
        disk.cleanup()
162 186

  
163
    #The image is ready, lets call kamaki if necessary
164
    if options.upload:
165
        pass
166

  
167
    if options.outfile is None:
168
        os.unlink(outfile)
169

  
170 187
    return 0
171 188

  
172
COLOR_BLACK = "\033[00m"
173
COLOR_RED = "\033[1;31m"
174

  
175 189

  
176 190
def main():
177 191
    try:
178 192
        ret = image_creator()
179 193
        sys.exit(ret)
180 194
    except FatalError as e:
181
        print >> sys.stderr, "\n%sError: %s%s\n" % (COLOR_RED, e, COLOR_BLACK)
195
        error(e)
182 196
        sys.exit(1)
183 197

  
184 198

  
b/image_creator/os_type/__init__.py
32 32
# or implied, of GRNET S.A.
33 33

  
34 34
import re
35
from clint.textui import indent, puts
35 36

  
36 37

  
37 38
def add_prefix(target):
......
111 112

  
112 113
    def data_cleanup(self):
113 114
        """Cleanup sensitive data out of the OS image."""
114
        raise NotImplementedError
115

  
116
        puts('Cleaning up sensitive data out of the OS image:')
117
        with indent(4):
118
            for name in dir(self):
119
                attr = getattr(self, name)
120
                if name.startswith('data_cleanup_') and callable(attr):
121
                    attr()
122
        puts()
115 123

  
116 124
    def sysprep(self):
117 125
        """Prepere system for image creation."""
118
        raise NotImplementedError
126

  
127
        puts('Preparing system for image creation:')
128
        with indent(4):
129
            for name in dir(self):
130
                attr = getattr(self, name)
131
                if name.startswith('sysprep_') and callable(attr):
132
                    attr()
133
        puts()
119 134

  
120 135
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
b/image_creator/os_type/linux.py
32 32
# or implied, of GRNET S.A.
33 33

  
34 34
from image_creator.os_type.unix import Unix
35
from image_creator.util import warn
36

  
37
from clint.textui import puts, indent
38

  
35 39
import re
36 40

  
37 41

  
......
53 57
                self._uuid[dev] = attr[1]
54 58
                return attr[1]
55 59

  
56
    def sysprep(self):
57
        """Prepere system for image creation."""
58
        self.sysprep_acpid()
59
        self.sysprep_persistent_net_rules()
60
        self.sysprep_persistent_devs()
61

  
62 60
    def sysprep_acpid(self):
63 61
        """Replace acpid powerdown action scripts to automatically shutdown
64 62
        the system without checking if a GUI is running.
65 63
        """
64

  
65
        puts('* Fixing acpid powerdown action')
66

  
66 67
        action = '#!/bin/sh\n\nPATH=/sbin:/bin:/usr/bin\n shutdown -h now '
67 68
        '\"Power button pressed\"'
68 69

  
......
71 72
        elif self.g.is_file('/etc/acpi/actions/power.sh'):
72 73
            self.g.write('/etc/acpi/actions/power.sh', action)
73 74
        else:
74
            print "Warning: No acpid action file found"
75
            with indent(2):
76
                warn("No acpid action file found")
75 77

  
76 78
    def sysprep_persistent_net_rules(self):
77 79
        """Remove udev rules that will keep network interface names persistent
78 80
        after hardware changes and reboots. Those rules will be created again
79 81
        the next time the image runs.
80 82
        """
83

  
84
        puts('* Removing persistent network interface names')
85

  
81 86
        rule_file = '/etc/udev/rules.d/70-persistent-net.rules'
82 87
        if self.g.is_file(rule_file):
83 88
            self.g.rm(rule_file)
......
86 91
        """Scan fstab and grub configuration files and replace all
87 92
        non-persistent device appearences with UUIDs.
88 93
        """
94

  
95
        puts('* Replacing fstab & grub non-persistent device appearences')
96

  
89 97
        # convert all devices in fstab to persistent
90 98
        persistent_root = self._persistent_fstab()
91 99

  
......
146 154

  
147 155
        entry = line.split()
148 156
        if len(entry) != 6:
149
            print "Warning: detected abnorman entry in fstab"
157
            warn("Detected abnormal entry in fstab")
150 158
            return orig, "", ""
151 159

  
152 160
        dev = entry[0]
b/image_creator/os_type/slackware.py
35 35

  
36 36

  
37 37
class Slackware(Linux):
38
    def cleanup_log(self):
38
    def data_cleanup_log(self):
39 39
        # In slackware the metadata about installed packages are
40 40
        # stored in /var/log/packages. Clearing all /var/log files
41 41
        # will destroy the package management system.
b/image_creator/os_type/unix.py
35 35
import sys
36 36

  
37 37
from image_creator.os_type import OSBase
38
from image_creator.util import warn
39
from clint.textui import puts
38 40

  
39 41

  
40 42
class Unix(OSBase):
......
63 65

  
64 66
            user, passwd = match.groups()
65 67
            if len(passwd) > 0 and passwd[0] == '!':
66
                print "Warning: Ignoring locked %s account." % user
68
                warn("Ignoring locked %s account." % user)
67 69
            else:
68 70
                users.append(user)
69 71

  
70 72
        return users
71 73

  
72
    def data_cleanup(self):
73
        self.data_cleanup_userdata()
74
        self.data_cleanup_tmp()
75
        self.data_cleanup_log()
76
        self.data_cleanup_mail()
77
        self.data_cleanup_cache()
78

  
79 74
    def data_cleanup_cache(self):
80 75
        """Remove all regular files under /var/cache"""
76

  
77
        puts('* Removing files under /var/cache')
78

  
81 79
        self.foreach_file('/var/cache', self.g.rm, ftype='r')
82 80

  
83 81
    def data_cleanup_tmp(self):
84 82
        """Remove all files under /tmp and /var/tmp"""
83

  
84
        puts('* Removing files under /tmp and /var/tmp')
85

  
85 86
        self.foreach_file('/tmp', self.g.rm_rf, maxdepth=1)
86 87
        self.foreach_file('/var/tmp', self.g.rm_rf, maxdepth=1)
87 88

  
88 89
    def data_cleanup_log(self):
89 90
        """Empty all files under /var/log"""
91

  
92
        puts('* Emptying all files under /var/log')
93

  
90 94
        self.foreach_file('/var/log', self.g.truncate, ftype='r')
91 95

  
92 96
    def data_cleanup_mail(self):
93 97
        """Remove all files under /var/mail and /var/spool/mail"""
98

  
99
        puts('* Removing files under /var/mail and /var/spool/mail')
100

  
94 101
        self.foreach_file('/var/spool/mail', self.g.rm_rf, maxdepth=1)
95 102
        self.foreach_file('/var/mail', self.g.rm_rf, maxdepth=1)
96 103

  
97 104
    def data_cleanup_userdata(self):
98 105
        """Delete sensitive userdata"""
106

  
99 107
        homedirs = ['/root'] + self.ls('/home/')
100 108

  
101 109
        for homedir in homedirs:
110
            puts('* Removing sensitive user data under %s' % homedir)
102 111
            for data in self.sensitive_userdata:
103 112
                fname = "%s/%s" % (homedir, data)
104 113
                if self.g.is_file(fname):
b/image_creator/util.py
32 32
# or implied, of GRNET S.A.
33 33

  
34 34
import pbs
35
from clint.textui import puts, puts_err, colored, progress
35 36

  
36 37

  
37 38
def get_command(command):
......
46 47
        return pbs.__getattr__(command)
47 48
    except pbs.CommadNotFount as e:
48 49
        return find_sbin_command(command, e)
50

  
51

  
52
def error(msg):
53
    puts_err(colored.red("Error: %s\n" % msg))
54

  
55

  
56
def warn(msg):
57
    puts_err(colored.yellow("Warning: %s" % msg))
58

  
59

  
60
def success(msg):
61
    puts(colored.green(msg))
62

  
63

  
64
def progress_generator(label='', n=100):
65
    position = 0
66
    for i in progress.bar(range(n), label):
67
        if i < position:
68
            continue
69
        position = yield
70
    yield  # suppress the StopIteration exception
71

  
72

  
73
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
b/setup.py
47 47
    license='BSD',
48 48
    packages=['image_creator'],
49 49
    include_package_data=True,
50
    install_requires=['pbs', 'clint'],
50
    install_requires=['pbs', 'clint', 'pysendfile'],
51 51
    entry_points={
52 52
        'console_scripts': ['snf-image-creator = image_creator.main:main']
53 53
    }

Also available in: Unified diff