From 81c0d1b3ad76e46bb76d4319ca93fc751d607dd5 Mon Sep 17 00:00:00 2001
From: Shuangquan Zhou <shuangquan.zhou@intel.com>
Date: Mon, 15 Aug 2011 14:42:46 +0800
Subject: [PATCH] Updated chroot's cleanup function

---
 plugins/imager/livecd_plugin.py  |  7 ++++---
 plugins/imager/liveusb_plugin.py |  8 +++++---
 plugins/imager/loop_plugin.py    |  6 +++---
 plugins/imager/raw_plugin.py     | 15 ++++++---------
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/plugins/imager/livecd_plugin.py b/plugins/imager/livecd_plugin.py
index f621bb8..c98a0cc 100644
--- a/plugins/imager/livecd_plugin.py
+++ b/plugins/imager/livecd_plugin.py
@@ -126,10 +126,11 @@ class LiveCDPlugin(ImagerPlugin):
         try:
             chroot.chroot(extmnt, None,  "/bin/env HOME=/root /bin/bash")
         except:
-            print >> sys.stderr, "Failed to chroot to %s." % img
+            raise CreatorError("Failed to chroot to %s." %img)
         finally:
-            chroot.cleanup_after_chroot("img",extloop,None,None)
-            return 1
+            imgloop.cleanup()
+            shutil.rmtree(imgmnt, ignore_errors = True) 
+            chroot.cleanup_after_chroot("img",extloop,tmpoutdir,extmnt)
         
     @classmethod
     def do_pack(cls, base_on):
diff --git a/plugins/imager/liveusb_plugin.py b/plugins/imager/liveusb_plugin.py
index c54c5dc..33af1a6 100644
--- a/plugins/imager/liveusb_plugin.py
+++ b/plugins/imager/liveusb_plugin.py
@@ -128,9 +128,11 @@ class LiveUSBPlugin(ImagerPlugin):
         try:
             chroot.chroot(extmnt, None,  "/bin/env HOME=/root /bin/bash")
         except:
-            chroot.cleanup_after_chroot("img", extloop, None, None)
-            print >> sys.stderr, "Failed to chroot to %s." % img
-            return 1
+            raise CreatorError("Failed to chroot to %s." %img)
+        finally:
+            imgloop.cleanup()
+            shutil.rmtree(imgmnt, ignore_errors = True) 
+            chroot.cleanup_after_chroot("img", extloop, tmpoutdir, extmnt)
     
     @classmethod
     def do_pack(cls, base_on):
diff --git a/plugins/imager/loop_plugin.py b/plugins/imager/loop_plugin.py
index ba02686..f743b6b 100644
--- a/plugins/imager/loop_plugin.py
+++ b/plugins/imager/loop_plugin.py
@@ -82,9 +82,9 @@ class LoopPlugin(ImagerPlugin):
         try:
             chroot.chroot(extmnt, None,  "/bin/env HOME=/root /bin/bash")
         except:
-            chroot.cleanup_after_chroot("img", extloop, None, None)
-            print >> sys.stderr, "Failed to chroot to %s." % extloop
-            return 1
+            raise CreatorError("Failed to chroot to %s." %img)  
+        finally:
+            chroot.cleanup_after_chroot("img", extloop, None, extmnt)
 
     @classmethod
     def do_unpack(cls, srcimg):
diff --git a/plugins/imager/raw_plugin.py b/plugins/imager/raw_plugin.py
index 8c7c140..bdb9076 100644
--- a/plugins/imager/raw_plugin.py
+++ b/plugins/imager/raw_plugin.py
@@ -71,11 +71,9 @@ class RawPlugin(ImagerPlugin):
         imgsize = misc.get_file_size(img) * 1024L * 1024L
         partedcmd = fs_related.find_binary_path("parted")
         disk = fs_related.SparseLoopbackDisk(img, imgsize)
-        extmnt = misc.mkdtemp()
-        tmpoutdir = misc.mkdtemp()
-        imgloop = PartitionedMount({'/dev/sdb':disk}, extmnt, skipformat = True)
+        imgmnt = misc.mkdtemp()
+        imgloop = PartitionedMount({'/dev/sdb':disk}, imgmnt, skipformat = True)
         img_fstype = "ext3"
-        extloop = None
         
         # Check the partitions from raw disk.
         p1 = subprocess.Popen([partedcmd,"-s",img,"unit","B","print"],
@@ -143,18 +141,17 @@ class RawPlugin(ImagerPlugin):
         
         try:
             imgloop.mount()
-            os_image = img
         except MountError, e:
             imgloop.cleanup()
             raise CreatorError("Failed to loopback mount '%s' : %s" %
                                (img, e))
 
         try:
-            chroot.chroot(extmnt, None,  "/bin/env HOME=/root /bin/bash")
+            chroot.chroot(imgmnt, None,  "/bin/env HOME=/root /bin/bash")
         except:
-            chroot.cleanup_after_chroot("img", imgloop, None, None)
-            print >> sys.stderr, "Failed to chroot to %s." % img
-            return 1
+            raise CreatorError("Failed to chroot to %s." %img)  
+        finally:
+            chroot.cleanup_after_chroot("img", imgloop, None, imgmnt)
             
     @classmethod
     def do_unpack(cls, srcimg):
-- 
2.7.4