Major code revision

* Add foreach_file function in OSBase. This is used by all cleanup
  methods
* Add slackware class. This class derives from linux and overwrites the
  log cleanup behaviour. This is needed since slackware stores all
  package management info in /var/log/packages
* Add a shrink method in DiskDevice. This method shrinks the last
  filesystem to the minimum size
* Make the code style comply with pep8

image_creator/__init__.py

@@ -35,6 +35,7 @@ __version__ = '0.1'
 
 import image_creator.os_type
 
+
 def get_os_class(distro, osfamily):
     module = None
     classname = None

image_creator/disk.py

@@ -13,7 +13,10 @@ from pbs import dmsetup
 from pbs import blockdev
 from pbs import dd
 
-class DiskError(Exception): pass
+
+class DiskError(Exception):
+    pass
+
 
 class Disk(object):
 
@@ -38,7 +41,7 @@ class Disk(object):
         while len(self._devices):
             device = self._devices.pop()
             device.destroy()
-            
+
         while len(self._cleanup_jobs):
             job, args = self._cleanup_jobs.pop()
             job(*args)
@@ -57,8 +60,8 @@ class Disk(object):
         size = blockdev('--getsize', sourcedev)
         cowfd, cow = tempfile.mkstemp()
         self._add_cleanup(os.unlink, cow)
-        # Create 1G cow file
-        dd('if=/dev/null', 'of=%s' % cow, 'bs=1k' ,'seek=%d' % (1024*1024))
+        # 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
@@ -79,9 +82,10 @@ class Disk(object):
         self._devices.remove(device)
         device.destroy()
 
+
 class DiskDevice(object):
 
-    def __init__(self, device, bootable = True):
+    def __init__(self, device, bootable=True):
         self.device = device
         self.bootable = bootable
 
@@ -89,7 +93,7 @@ class DiskDevice(object):
 
         self.g.set_trace(1)
 
-        self.g.add_drive_opts(device, readonly = 0)
+        self.g.add_drive_opts(device, readonly=0)
         self.g.launch()
         roots = self.g.inspect_os()
         if len(roots) == 0:
@@ -100,22 +104,26 @@ class DiskDevice(object):
         self.root = roots[0]
         self.ostype = self.g.inspect_get_type(self.root)
         self.distro = self.g.inspect_get_distro(self.root)
-    
+
     def destroy(self):
         self.g.umount_all()
         self.g.sync()
         # Close the guestfs handler
         self.g.close()
         del self.g
-    
+
     def mount(self):
         mps = self.g.inspect_get_mountpoints(self.root)
+
         # Sort the keys to mount the fs in a correct order.
         # / should be mounted befor /boot, etc
-        def compare (a, b):
-            if len(a[0]) > len(b[0]): return 1
-            elif len(a[0]) == len(b[0]): return 0
-            else: return -1
+        def compare(a, b):
+            if len(a[0]) > len(b[0]):
+                return 1
+            elif len(a[0]) == len(b[0]):
+                return 0
+            else:
+                return -1
         mps.sort(compare)
         for mp, dev in mps:
             try:
@@ -123,4 +131,39 @@ class DiskDevice(object):
             except RuntimeError as msg:
                 print "%s (ignored)" % msg
 
+    def umount(self):
+        self.g.umount_all()
+
+    def shrink(self):
+        dev = self.g.part_to_dev(self.root)
+        parttype = self.g.part_get_parttype(dev)
+        if parttype != 'msdos':
+            raise DiskError("You have a %s partition table. "
+                "Only msdos partitions are supported" % parttype)
+
+        last_partition = self.g.part_list(dev)[-1]
+
+        if last_partition['part_num'] > 4:
+            raise DiskError("This disk contains logical partitions. "
+                "Only primary partitions are supported.")
+
+        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
+            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])
+
+        sector_size = self.g.blockdev_getss(dev)
+
+        start = last_partition['part_start'] / sector_size
+        end = start + (block_size * block_cnt) / sector_size - 1
+
+
+
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/main.py

@@ -36,6 +36,7 @@ from image_creator.disk import Disk
 import sys
 import os
 
+
 def main():
     if len(sys.argv) != 3:
         sys.exit("Usage: %s <source> <output_file>" %
@@ -50,10 +51,9 @@ def main():
         osclass = get_os_class(dev.distro, dev.ostype)
         image_os = osclass(dev.root, dev.g)
         metadata = image_os.get_metadata()
-        for key, val in metadata.iteritems():
-            print "%s=%s" % (key,val)
-
         image_os.data_cleanup()
+        dev.umount()
+        dev.shrink()
 
     finally:
         disk.cleanup()
@@ -62,4 +62,3 @@ if __name__ == '__main__':
     main()
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
-

image_creator/os_type/__init__.py

 #!/usr/bin/env python
 
+import re
+
+
 def add_prefix(target):
     def wrapper(self, *args):
         prefix = args[0]
         return map(lambda x: prefix + x, target(self, *args))
     return wrapper
 
+
 class OSBase(object):
     def __init__(self, rootdev, ghandler):
         self.root = rootdev
         self.g = ghandler
 
     @add_prefix
-    def ls(self, directory): return self.g.ls(directory)
+    def ls(self, directory):
+        return self.g.ls(directory)
 
     @add_prefix
-    def find(self, directory): return self.g.find(directory)
+    def find(self, directory):
+        return self.g.find(directory)
+
+    def foreach_file(self, directory, action, **kargs):
+
+        maxdepth = None if 'maxdepth' not in kargs else kargs['maxdepth']
+        if maxdepth == 0:
+            return
+
+        # maxdepth -= 1
+        maxdepth = None if maxdepth is None else maxdepth - 1
+        kargs['maxdepth'] = maxdepth
+
+        exclude = None if 'exclude' not in kargs else kargs['exclude']
+        ftype = None if 'ftype' not in kargs else kargs['ftype']
+        has_ftype = lambda x, y: y is None and True or x['ftyp'] == y
+
+        for f in self.g.readdir(directory):
+            if f['name'] in ('.', '..'):
+                continue
+
+            full_path = "%s/%s" % (directory, f['name'])
+
+            if exclude and re.match(exclude, full_path):
+                continue
+
+            if has_ftype(f, 'd'):
+                self.foreach_file(full_path, action, **kargs)
+
+            if has_ftype(f, ftype):
+                action(full_path)
 
     def get_metadata(self):
         meta = {}
@@ -24,23 +59,8 @@ class OSBase(object):
         meta["description"] = self.g.inspect_get_product_name(self.root)
 
         return meta
-
-    def mount_all(self):
-        mps = g.inspect_get_mountpoints(self.root)
-        # Sort the keys to mount the fs in a correct order.
-        # / should be mounted befor /boot, etc
-        def compare (a, b):
-            if len(a[0]) > len(b[0]): return 1
-            elif len(a[0]) == len(b[0]): return 0
-            else: return -1
-        mps.sort(compare)
-        for mp, dev in mps:
-            try:
-                self.g.mount(dev, mp)
-            except RuntimeError as msg:
-                print "%s (ignored)" % msg
-
-    def cleanup_sensitive_data(self):
+        
+    def data_cleanup(self):
         raise NotImplementedError
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/os_type/freebsd.py

 from image_creator.os_type.unix import Unix
 
+
 class Freebsd(Unix):
-	pass
+    pass
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/os_type/hurd.py

 from image_creator.os_type.unix import Unix
 
+
 class Hard(Unix):
-	pass
+    pass
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/os_type/linux.py

 from image_creator.os_type.unix import Unix
 
+
 class Linux(Unix):
-	pass
+    pass
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/os_type/netbsd.py

 from image_creator.os_type.unix import Unix
 
+
 class Netbsd(Unix):
-	pass
+    pass
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/os_type/slackware.py

+from image_creator.os_type.linux import Linux
+
+
+class Slackware(Linux):
+    def cleanup_log(self):
+        # In slackware the the installed packages info are stored in
+        # /var/log/packages. Clearing all /var/log files will destroy
+        # the package management
+        self.foreach_file('/var/log', self.g.truncate, ftype='r', \
+            exclude='/var/log/packages')
+
+# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

image_creator/os_type/unix.py

 #!/usr/bin/env python
 
 import re
+import sys
 
 from image_creator.os_type import OSBase
 
+
 class Unix(OSBase):
 
     sensitive_userdata = ['.bash_history']
@@ -36,26 +38,17 @@ class Unix(OSBase):
         self.cleanup_log()
 
     def cleanup_tmp(self):
-        files = []
-        files.extend(self.ls('/tmp/'))
-        files.extend(self.ls('/var/tmp/'))
-    
-        for filename in files:
-            self.g.rm_rf(filename)
+        self.foreach_file('/tmp', self.g.rm_rf, maxdepth=1)
 
     def cleanup_log(self):
-        files = self.find( '/var/log/')
+        self.foreach_file('/var/log', self.g.truncate, ftype='r')
 
-        for filename in filter(self.g.is_file, files):
-            self.g.truncate(filename)
-        
     def cleanup_userdata(self):
         homedirs = ['/root'] + self.ls('/home/')
 
         for homedir in homedirs:
             for data in self.sensitive_userdata:
                 fname = "%s/%s" % (homedir, data)
-                print "Filename: %s\n" % fname
                 if self.g.is_file(fname):
                     self.g.scrub_file(fname)
 

image_creator/os_type/windows.py

 from image_creator.os_type import OSBase
 
+
 class Windows(OSBase):
-	pass
+    pass
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :