From: Nikos Skalkotos <skalkoto@grnet.gr>
Date: Tue, 20 Mar 2012 16:04:38 +0000 (+0200)
Subject: Beautify program's output.
X-Git-Tag: v0.1~128
X-Git-Url: https://code.grnet.gr/git/snf-image-creator/commitdiff_plain/22a6d23262e7b25844bf26dc5fc95ce33f4726b2

Beautify program's output.
---

diff --git a/image_creator/disk.py b/image_creator/disk.py
index 874d593..2213c86 100644
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -31,9 +31,9 @@
 # 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 clint.textui import progress
+from clint.textui import indent, puts, colored
 
 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.
         """
-        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
-        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()
@@ -139,15 +158,6 @@ class Disk(object):
         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.
@@ -170,27 +180,32 @@ class DiskDevice(object):
 
     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)
-        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."""
@@ -242,6 +257,8 @@ class DiskDevice(object):
         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':
@@ -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):
-            print "Warning: Don't know how to resize %s partitions." % vfs_type
+            warn("Don't know how to resize %s partitions." % vfs_type)
             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.
diff --git a/image_creator/main.py b/image_creator/main.py
index a225579..e4f9a29 100755
--- a/image_creator/main.py
+++ b/image_creator/main.py
@@ -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.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
@@ -106,8 +108,35 @@ def parse_options(input_args):
     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:
@@ -125,6 +154,7 @@ def image_creator():
     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()
@@ -140,45 +170,29 @@ def image_creator():
         size = options.shrink and dev.shrink() or dev.size()
         metadata['size'] = str(size // 2 ** 20)
 
-        outfile = ""
         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()
-        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:
+        puts('cleaning up...')
         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
 
-COLOR_BLACK = "\033[00m"
-COLOR_RED = "\033[1;31m"
-
 
 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)
 
 
diff --git a/image_creator/os_type/__init__.py b/image_creator/os_type/__init__.py
index ab14356..f3fd173 100644
--- a/image_creator/os_type/__init__.py
+++ b/image_creator/os_type/__init__.py
@@ -32,6 +32,7 @@
 # or implied, of GRNET S.A.
 
 import re
+from clint.textui import indent, puts
 
 
 def add_prefix(target):
@@ -111,10 +112,24 @@ class OSBase(object):
 
     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."""
-        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 :
diff --git a/image_creator/os_type/linux.py b/image_creator/os_type/linux.py
index 4f08332..066f6bd 100644
--- a/image_creator/os_type/linux.py
+++ b/image_creator/os_type/linux.py
@@ -32,6 +32,10 @@
 # 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
 
 
@@ -53,16 +57,13 @@ class Linux(Unix):
                 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.
         """
+
+        puts('* Fixing acpid powerdown action')
+
         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:
-            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.
         """
+
+        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)
@@ -86,6 +91,9 @@ class Linux(Unix):
         """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()
 
@@ -146,7 +154,7 @@ class Linux(Unix):
 
         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]
diff --git a/image_creator/os_type/slackware.py b/image_creator/os_type/slackware.py
index e52bf20..d4f4c44 100644
--- a/image_creator/os_type/slackware.py
+++ b/image_creator/os_type/slackware.py
@@ -35,7 +35,7 @@ from image_creator.os_type.linux import 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.
diff --git a/image_creator/os_type/unix.py b/image_creator/os_type/unix.py
index 20b971c..2096c7b 100644
--- a/image_creator/os_type/unix.py
+++ b/image_creator/os_type/unix.py
@@ -35,6 +35,8 @@ import re
 import sys
 
 from image_creator.os_type import OSBase
+from image_creator.util import warn
+from clint.textui import puts
 
 
 class Unix(OSBase):
@@ -63,42 +65,49 @@ class Unix(OSBase):
 
             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
 
-    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"""
+
+        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"""
+
+        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"""
+
+        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"""
+
+        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"""
+
         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):
diff --git a/image_creator/util.py b/image_creator/util.py
index f912fe7..bb57b5f 100644
--- a/image_creator/util.py
+++ b/image_creator/util.py
@@ -32,6 +32,7 @@
 # or implied, of GRNET S.A.
 
 import pbs
+from clint.textui import puts, puts_err, colored, progress
 
 
 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)
+
+
+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 :
diff --git a/setup.py b/setup.py
index 3e0384f..0e4d6b0 100755
--- a/setup.py
+++ b/setup.py
@@ -47,7 +47,7 @@ setup(
     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']
     }