Beautify program's output.
authorNikos Skalkotos <skalkoto@grnet.gr>
Tue, 20 Mar 2012 16:04:38 +0000 (18:04 +0200)
committerNikos Skalkotos <skalkoto@grnet.gr>
Tue, 20 Mar 2012 16:04:38 +0000 (18:04 +0200)
image_creator/disk.py
image_creator/main.py
image_creator/os_type/__init__.py
image_creator/os_type/linux.py
image_creator/os_type/slackware.py
image_creator/os_type/unix.py
image_creator/util.py
setup.py

index 874d593..2213c86 100644 (file)
@@ -31,9 +31,9 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.util import get_command
+from image_creator.util import get_command, warn, progress_generator
 from image_creator import FatalError
 from image_creator import FatalError
-from clint.textui import progress
+from clint.textui import indent, puts, colored
 
 import stat
 import os
 
 import stat
 import os
@@ -99,33 +99,52 @@ class Disk(object):
         This instance is a snapshot of the original source media of
         the Disk instance.
         """
         This instance is a snapshot of the original source media of
         the Disk instance.
         """
-        sourcedev = self.source
-        mode = os.stat(self.source).st_mode
-        if stat.S_ISDIR(mode):
-            return self._losetup(self._dir_to_disk())
-        elif stat.S_ISREG(mode):
-            sourcedev = self._losetup(self.source)
-        elif not stat.S_ISBLK(mode):
-            raise ValueError("Value for self.source is invalid")
+
+        puts("Examining source media `%s'" % self.source)
+        with indent(4):
+            sourcedev = self.source
+            mode = os.stat(self.source).st_mode
+            if stat.S_ISDIR(mode):
+                puts(colored.green('Looks like a directory'))
+                return self._losetup(self._dir_to_disk())
+            elif stat.S_ISREG(mode):
+                puts(colored.green('Looks like an image file'))
+                sourcedev = self._losetup(self.source)
+            elif not stat.S_ISBLK(mode):
+                raise ValueError("Invalid media source. Only block devices, "
+                                "regular files and directories are supported.")
+            else:
+                puts(colored.green('Looks like a block device'))
+            #puts()
 
         # Take a snapshot and return it to the user
 
         # Take a snapshot and return it to the user
-        size = blockdev('--getsize', sourcedev)
-        cowfd, cow = tempfile.mkstemp()
-        os.close(cowfd)
-        self._add_cleanup(os.unlink, cow)
-        # Create 1G cow sparse file
-        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
-        cowdev = self._losetup(cow)
-
-        snapshot = uuid.uuid4().hex
-        tablefd, table = tempfile.mkstemp()
-        try:
-            os.write(tablefd, "0 %d snapshot %s %s n 8" % \
-                                        (int(size), sourcedev, cowdev))
-            dmsetup('create', snapshot, table)
-            self._add_cleanup(dmsetup, 'remove', snapshot)
-        finally:
-            os.unlink(table)
+        puts("Snapshotting media source")
+        with indent(4):
+            size = blockdev('--getsize', sourcedev)
+            cowfd, cow = tempfile.mkstemp()
+            os.close(cowfd)
+            self._add_cleanup(os.unlink, cow)
+            # Create 1G cow sparse file
+            dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
+                                            'seek=%d' % (1024 * 1024))
+            cowdev = self._losetup(cow)
+    
+            snapshot = uuid.uuid4().hex
+            tablefd, table = tempfile.mkstemp()
+            try:
+                os.write(tablefd, "0 %d snapshot %s %s n 8" % \
+                                            (int(size), sourcedev, cowdev))
+                dmsetup('create', snapshot, table)
+                self._add_cleanup(dmsetup, 'remove', snapshot)
+                # Sometimes dmsetup remove fails with Device or resource busy,
+                # although everything is cleaned up and the snapshot is not
+                # used by anyone. Add a 2 seconds delay to be on the safe side.
+                self._add_cleanup(time.sleep, 2)
+
+            finally:
+                os.unlink(table)
+            puts(colored.green('Done'))
+        # puts()
         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
         self._devices.append(new_device)
         new_device.enable()
         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
         self._devices.append(new_device)
         new_device.enable()
@@ -139,15 +158,6 @@ class Disk(object):
         device.destroy()
 
 
         device.destroy()
 
 
-def progress_generator(label=''):
-    position = 0
-    for i in progress.bar(range(100), label):
-        if i < position:
-            continue
-        position = yield
-    yield  # suppress the StopIteration exception
-
-
 class DiskDevice(object):
     """This class represents a block device hosting an Operating System
     as created by the device-mapper.
 class DiskDevice(object):
     """This class represents a block device hosting an Operating System
     as created by the device-mapper.
@@ -170,27 +180,32 @@ class DiskDevice(object):
 
     def enable(self):
         """Enable a newly created DiskDevice"""
 
     def enable(self):
         """Enable a newly created DiskDevice"""
-
-        self.progressbar = progress_generator("VM lauch: ")
-        self.progressbar.next()
-        eh = self.g.set_event_callback(self.progress_callback,
+        self.progressbar = progress_generator("Launching helper VM: ")
+        with indent(4):
+            self.progressbar.next()
+            eh = self.g.set_event_callback(self.progress_callback,
                                                         guestfs.EVENT_PROGRESS)
                                                         guestfs.EVENT_PROGRESS)
-        self.g.launch()
-        self.guestfs_enabled = True
-        self.g.delete_event_callback(eh)
-        if self.progressbar is not None:
-            self.progressbar.send(100)
-            self.progressbar = None
-
-        roots = self.g.inspect_os()
-        if len(roots) == 0:
-            raise FatalError("No operating system found")
-        if len(roots) > 1:
-            raise FatalError("Multiple operating systems found")
-
-        self.root = roots[0]
-        self.ostype = self.g.inspect_get_type(self.root)
-        self.distro = self.g.inspect_get_distro(self.root)
+            self.g.launch()
+            self.guestfs_enabled = True
+            self.g.delete_event_callback(eh)
+            if self.progressbar is not None:
+                self.progressbar.send(100)
+                self.progressbar = None
+            puts(colored.green('Done'))
+
+        puts('Inspecting Operating System')
+        with indent(4):
+            roots = self.g.inspect_os()
+            if len(roots) == 0:
+                raise FatalError("No operating system found")
+            if len(roots) > 1:
+                raise FatalError("Multiple operating systems found."
+                                "We only support images with one filesystem.")
+            self.root = roots[0]
+            self.ostype = self.g.inspect_get_type(self.root)
+            self.distro = self.g.inspect_get_distro(self.root)
+            puts(colored.green('Found a %s system' % self.distro))
+        puts()
 
     def destroy(self):
         """Destroy this DiskDevice instance."""
 
     def destroy(self):
         """Destroy this DiskDevice instance."""
@@ -242,6 +257,8 @@ class DiskDevice(object):
         disk and then updating the partition table. The new disk size
         (in bytes) is returned.
         """
         disk and then updating the partition table. The new disk size
         (in bytes) is returned.
         """
+        puts("Shrinking image (this may take a while)")
+        
         dev = self.g.part_to_dev(self.root)
         parttype = self.g.part_get_parttype(dev)
         if parttype != 'msdos':
         dev = self.g.part_to_dev(self.root)
         parttype = self.g.part_get_parttype(dev)
         if parttype != 'msdos':
@@ -257,24 +274,30 @@ class DiskDevice(object):
         part_dev = "%s%d" % (dev, last_partition['part_num'])
         fs_type = self.g.vfs_type(part_dev)
         if not re.match("ext[234]", fs_type):
         part_dev = "%s%d" % (dev, last_partition['part_num'])
         fs_type = self.g.vfs_type(part_dev)
         if not re.match("ext[234]", fs_type):
-            print "Warning: Don't know how to resize %s partitions." % vfs_type
+            warn("Don't know how to resize %s partitions." % vfs_type)
             return
 
             return
 
-        self.g.e2fsck_f(part_dev)
-        self.g.resize2fs_M(part_dev)
-        output = self.g.tune2fs_l(part_dev)
-        block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
-        block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
+        with indent(4):
+            self.g.e2fsck_f(part_dev)
+            self.g.resize2fs_M(part_dev)
+
+            output = self.g.tune2fs_l(part_dev)
+            block_size = int(
+                filter(lambda x: x[0] == 'Block size', output)[0][1])
+            block_cnt = int(
+                filter(lambda x: x[0] == 'Block count', output)[0][1])
 
 
-        sector_size = self.g.blockdev_getss(dev)
+            sector_size = self.g.blockdev_getss(dev)
 
 
-        start = last_partition['part_start'] / sector_size
-        end = start + (block_size * block_cnt) / sector_size - 1
+            start = last_partition['part_start'] / sector_size
+            end = start + (block_size * block_cnt) / sector_size - 1
 
 
-        self.g.part_del(dev, last_partition['part_num'])
-        self.g.part_add(dev, 'p', start, end)
+            self.g.part_del(dev, last_partition['part_num'])
+            self.g.part_add(dev, 'p', start, end)
 
 
-        return (end + 1) * sector_size
+            new_size = (end + 1) * sector_size
+            puts("  New image size is %dMB\n" % (new_size // 2 ** 20))
+        return new_size
 
     def size(self):
         """Returns the "payload" size of the device.
 
     def size(self):
         """Returns the "payload" size of the device.
index a225579..e4f9a29 100755 (executable)
@@ -37,7 +37,9 @@ from image_creator import get_os_class
 from image_creator import __version__ as version
 from image_creator import FatalError
 from image_creator.disk import Disk
 from image_creator import __version__ as version
 from image_creator import FatalError
 from image_creator.disk import Disk
-from image_creator.util import get_command
+from image_creator.util import get_command, error, progress_generator, success
+from clint.textui import puts, indent
+from sendfile import sendfile
 
 import sys
 import os
 
 import sys
 import os
@@ -106,8 +108,35 @@ def parse_options(input_args):
     return options
 
 
     return options
 
 
-def image_creator():
+def extract_image(device, outfile, size):
+    blocksize = 4194304  # 4MB
+    progress_size = (size + 1048575) // 1048576  # in MB
+    progressbar = progress_generator("Dumping image file: ",
+                                                    progress_size)
+    source = open(device, "r")
+    try:
+        dest = open(outfile, "w")
+        try:
+            left = size
+            offset = 0
+            progressbar.next()
+            while left > 0:
+                length = min(left, blocksize)
+                sent = sendfile(dest.fileno(), source.fileno(), offset, length)
+                offset += sent
+                left -= sent
+                for i in range(4):
+                    progressbar.next()
+        finally:
+            dest.close()
+    finally:
+        source.close()
+
+    success('Image file %s was successfully created' % outfile)
+
 
 
+def image_creator():
+    puts('snf-image-creator %s\n' % version)
     options = parse_options(sys.argv[1:])
 
     if os.geteuid() != 0:
     options = parse_options(sys.argv[1:])
 
     if os.geteuid() != 0:
@@ -125,6 +154,7 @@ def image_creator():
     try:
         dev = disk.get_device()
         dev.mount()
     try:
         dev = disk.get_device()
         dev.mount()
+
         osclass = get_os_class(dev.distro, dev.ostype)
         image_os = osclass(dev.root, dev.g)
         metadata = image_os.get_metadata()
         osclass = get_os_class(dev.distro, dev.ostype)
         image_os = osclass(dev.root, dev.g)
         metadata = image_os.get_metadata()
@@ -140,45 +170,29 @@ def image_creator():
         size = options.shrink and dev.shrink() or dev.size()
         metadata['size'] = str(size // 2 ** 20)
 
         size = options.shrink and dev.shrink() or dev.size()
         metadata['size'] = str(size // 2 ** 20)
 
-        outfile = ""
         if options.outfile is not None:
         if options.outfile is not None:
-            outfile = options.outfile
             f = open('%s.%s' % (options.outfile, 'meta'), 'w')
             try:
                 for key in metadata.keys():
                     f.write("%s=%s\n" % (key, metadata[key]))
             finally:
                 f.close()
             f = open('%s.%s' % (options.outfile, 'meta'), 'w')
             try:
                 for key in metadata.keys():
                     f.write("%s=%s\n" % (key, metadata[key]))
             finally:
                 f.close()
-        else:
-            outfd, outfile = tmpfile.mkstemp()
-            os.close(outfd)
 
 
-        dd('if=%s' % dev.device,
-            'of=%s' % outfile,
-            'bs=4M', 'count=%d' % ((size + 1) // 2 ** 22))
+            extract_image(dev.device, options.outfile, size)
 
     finally:
 
     finally:
+        puts('cleaning up...')
         disk.cleanup()
 
         disk.cleanup()
 
-    #The image is ready, lets call kamaki if necessary
-    if options.upload:
-        pass
-
-    if options.outfile is None:
-        os.unlink(outfile)
-
     return 0
 
     return 0
 
-COLOR_BLACK = "\033[00m"
-COLOR_RED = "\033[1;31m"
-
 
 def main():
     try:
         ret = image_creator()
         sys.exit(ret)
     except FatalError as e:
 
 def main():
     try:
         ret = image_creator()
         sys.exit(ret)
     except FatalError as e:
-        print >> sys.stderr, "\n%sError: %s%s\n" % (COLOR_RED, e, COLOR_BLACK)
+        error(e)
         sys.exit(1)
 
 
         sys.exit(1)
 
 
index ab14356..f3fd173 100644 (file)
@@ -32,6 +32,7 @@
 # or implied, of GRNET S.A.
 
 import re
 # or implied, of GRNET S.A.
 
 import re
+from clint.textui import indent, puts
 
 
 def add_prefix(target):
 
 
 def add_prefix(target):
@@ -111,10 +112,24 @@ class OSBase(object):
 
     def data_cleanup(self):
         """Cleanup sensitive data out of the OS image."""
 
     def data_cleanup(self):
         """Cleanup sensitive data out of the OS image."""
-        raise NotImplementedError
+
+        puts('Cleaning up sensitive data out of the OS image:')
+        with indent(4):
+            for name in dir(self):
+                attr = getattr(self, name)
+                if name.startswith('data_cleanup_') and callable(attr):
+                    attr()
+        puts()
 
     def sysprep(self):
         """Prepere system for image creation."""
 
     def sysprep(self):
         """Prepere system for image creation."""
-        raise NotImplementedError
+
+        puts('Preparing system for image creation:')
+        with indent(4):
+            for name in dir(self):
+                attr = getattr(self, name)
+                if name.startswith('sysprep_') and callable(attr):
+                    attr()
+        puts()
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
index 4f08332..066f6bd 100644 (file)
 # or implied, of GRNET S.A.
 
 from image_creator.os_type.unix import Unix
 # or implied, of GRNET S.A.
 
 from image_creator.os_type.unix import Unix
+from image_creator.util import warn
+
+from clint.textui import puts, indent
+
 import re
 
 
 import re
 
 
@@ -53,16 +57,13 @@ class Linux(Unix):
                 self._uuid[dev] = attr[1]
                 return attr[1]
 
                 self._uuid[dev] = attr[1]
                 return attr[1]
 
-    def sysprep(self):
-        """Prepere system for image creation."""
-        self.sysprep_acpid()
-        self.sysprep_persistent_net_rules()
-        self.sysprep_persistent_devs()
-
     def sysprep_acpid(self):
         """Replace acpid powerdown action scripts to automatically shutdown
         the system without checking if a GUI is running.
         """
     def sysprep_acpid(self):
         """Replace acpid powerdown action scripts to automatically shutdown
         the system without checking if a GUI is running.
         """
+
+        puts('* Fixing acpid powerdown action')
+
         action = '#!/bin/sh\n\nPATH=/sbin:/bin:/usr/bin\n shutdown -h now '
         '\"Power button pressed\"'
 
         action = '#!/bin/sh\n\nPATH=/sbin:/bin:/usr/bin\n shutdown -h now '
         '\"Power button pressed\"'
 
@@ -71,13 +72,17 @@ class Linux(Unix):
         elif self.g.is_file('/etc/acpi/actions/power.sh'):
             self.g.write('/etc/acpi/actions/power.sh', action)
         else:
         elif self.g.is_file('/etc/acpi/actions/power.sh'):
             self.g.write('/etc/acpi/actions/power.sh', action)
         else:
-            print "Warning: No acpid action file found"
+            with indent(2):
+                warn("No acpid action file found")
 
     def sysprep_persistent_net_rules(self):
         """Remove udev rules that will keep network interface names persistent
         after hardware changes and reboots. Those rules will be created again
         the next time the image runs.
         """
 
     def sysprep_persistent_net_rules(self):
         """Remove udev rules that will keep network interface names persistent
         after hardware changes and reboots. Those rules will be created again
         the next time the image runs.
         """
+
+        puts('* Removing persistent network interface names')
+
         rule_file = '/etc/udev/rules.d/70-persistent-net.rules'
         if self.g.is_file(rule_file):
             self.g.rm(rule_file)
         rule_file = '/etc/udev/rules.d/70-persistent-net.rules'
         if self.g.is_file(rule_file):
             self.g.rm(rule_file)
@@ -86,6 +91,9 @@ class Linux(Unix):
         """Scan fstab and grub configuration files and replace all
         non-persistent device appearences with UUIDs.
         """
         """Scan fstab and grub configuration files and replace all
         non-persistent device appearences with UUIDs.
         """
+
+        puts('* Replacing fstab & grub non-persistent device appearences')
+
         # convert all devices in fstab to persistent
         persistent_root = self._persistent_fstab()
 
         # convert all devices in fstab to persistent
         persistent_root = self._persistent_fstab()
 
@@ -146,7 +154,7 @@ class Linux(Unix):
 
         entry = line.split()
         if len(entry) != 6:
 
         entry = line.split()
         if len(entry) != 6:
-            print "Warning: detected abnorman entry in fstab"
+            warn("Detected abnormal entry in fstab")
             return orig, "", ""
 
         dev = entry[0]
             return orig, "", ""
 
         dev = entry[0]
index e52bf20..d4f4c44 100644 (file)
@@ -35,7 +35,7 @@ from image_creator.os_type.linux import Linux
 
 
 class Slackware(Linux):
 
 
 class Slackware(Linux):
-    def cleanup_log(self):
+    def data_cleanup_log(self):
         # In slackware the metadata about installed packages are
         # stored in /var/log/packages. Clearing all /var/log files
         # will destroy the package management system.
         # In slackware the metadata about installed packages are
         # stored in /var/log/packages. Clearing all /var/log files
         # will destroy the package management system.
index 20b971c..2096c7b 100644 (file)
@@ -35,6 +35,8 @@ import re
 import sys
 
 from image_creator.os_type import OSBase
 import sys
 
 from image_creator.os_type import OSBase
+from image_creator.util import warn
+from clint.textui import puts
 
 
 class Unix(OSBase):
 
 
 class Unix(OSBase):
@@ -63,42 +65,49 @@ class Unix(OSBase):
 
             user, passwd = match.groups()
             if len(passwd) > 0 and passwd[0] == '!':
 
             user, passwd = match.groups()
             if len(passwd) > 0 and passwd[0] == '!':
-                print "Warning: Ignoring locked %s account." % user
+                warn("Ignoring locked %s account." % user)
             else:
                 users.append(user)
 
         return users
 
             else:
                 users.append(user)
 
         return users
 
-    def data_cleanup(self):
-        self.data_cleanup_userdata()
-        self.data_cleanup_tmp()
-        self.data_cleanup_log()
-        self.data_cleanup_mail()
-        self.data_cleanup_cache()
-
     def data_cleanup_cache(self):
         """Remove all regular files under /var/cache"""
     def data_cleanup_cache(self):
         """Remove all regular files under /var/cache"""
+
+        puts('* Removing files under /var/cache')
+
         self.foreach_file('/var/cache', self.g.rm, ftype='r')
 
     def data_cleanup_tmp(self):
         """Remove all files under /tmp and /var/tmp"""
         self.foreach_file('/var/cache', self.g.rm, ftype='r')
 
     def data_cleanup_tmp(self):
         """Remove all files under /tmp and /var/tmp"""
+
+        puts('* Removing files under /tmp and /var/tmp')
+
         self.foreach_file('/tmp', self.g.rm_rf, maxdepth=1)
         self.foreach_file('/var/tmp', self.g.rm_rf, maxdepth=1)
 
     def data_cleanup_log(self):
         """Empty all files under /var/log"""
         self.foreach_file('/tmp', self.g.rm_rf, maxdepth=1)
         self.foreach_file('/var/tmp', self.g.rm_rf, maxdepth=1)
 
     def data_cleanup_log(self):
         """Empty all files under /var/log"""
+
+        puts('* Emptying all files under /var/log')
+
         self.foreach_file('/var/log', self.g.truncate, ftype='r')
 
     def data_cleanup_mail(self):
         """Remove all files under /var/mail and /var/spool/mail"""
         self.foreach_file('/var/log', self.g.truncate, ftype='r')
 
     def data_cleanup_mail(self):
         """Remove all files under /var/mail and /var/spool/mail"""
+
+        puts('* Removing files under /var/mail and /var/spool/mail')
+
         self.foreach_file('/var/spool/mail', self.g.rm_rf, maxdepth=1)
         self.foreach_file('/var/mail', self.g.rm_rf, maxdepth=1)
 
     def data_cleanup_userdata(self):
         """Delete sensitive userdata"""
         self.foreach_file('/var/spool/mail', self.g.rm_rf, maxdepth=1)
         self.foreach_file('/var/mail', self.g.rm_rf, maxdepth=1)
 
     def data_cleanup_userdata(self):
         """Delete sensitive userdata"""
+
         homedirs = ['/root'] + self.ls('/home/')
 
         for homedir in homedirs:
         homedirs = ['/root'] + self.ls('/home/')
 
         for homedir in homedirs:
+            puts('* Removing sensitive user data under %s' % homedir)
             for data in self.sensitive_userdata:
                 fname = "%s/%s" % (homedir, data)
                 if self.g.is_file(fname):
             for data in self.sensitive_userdata:
                 fname = "%s/%s" % (homedir, data)
                 if self.g.is_file(fname):
index f912fe7..bb57b5f 100644 (file)
@@ -32,6 +32,7 @@
 # or implied, of GRNET S.A.
 
 import pbs
 # or implied, of GRNET S.A.
 
 import pbs
+from clint.textui import puts, puts_err, colored, progress
 
 
 def get_command(command):
 
 
 def get_command(command):
@@ -46,3 +47,27 @@ def get_command(command):
         return pbs.__getattr__(command)
     except pbs.CommadNotFount as e:
         return find_sbin_command(command, e)
         return pbs.__getattr__(command)
     except pbs.CommadNotFount as e:
         return find_sbin_command(command, e)
+
+
+def error(msg):
+    puts_err(colored.red("Error: %s\n" % msg))
+
+
+def warn(msg):
+    puts_err(colored.yellow("Warning: %s" % msg))
+
+
+def success(msg):
+    puts(colored.green(msg))
+
+
+def progress_generator(label='', n=100):
+    position = 0
+    for i in progress.bar(range(n), label):
+        if i < position:
+            continue
+        position = yield
+    yield  # suppress the StopIteration exception
+
+
+# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
index 3e0384f..0e4d6b0 100755 (executable)
--- a/setup.py
+++ b/setup.py
@@ -47,7 +47,7 @@ setup(
     license='BSD',
     packages=['image_creator'],
     include_package_data=True,
     license='BSD',
     packages=['image_creator'],
     include_package_data=True,
-    install_requires=['pbs', 'clint'],
+    install_requires=['pbs', 'clint', 'pysendfile'],
     entry_points={
         'console_scripts': ['snf-image-creator = image_creator.main:main']
     }
     entry_points={
         'console_scripts': ['snf-image-creator = image_creator.main:main']
     }