Expand to allow operation on a rootfs directory.

1 files changed, 42 insertions, 17 deletions

diff --git a/vmsquashtar.py b/vmsquashtar.py
index c8bbae0..2cf888e 100755
--- a/vmsquashtar.py
+++ b/vmsquashtar.py
@@ -30,7 +30,7 @@ from vmdebootstrap.filesystem import Filesystem
 
 
 __version__ = '0.1'
-__desc__ = "Helper to make a squashfs from a tarball."
+__desc__ = "Helper to make a squashfs from a tarball or directory."
 
 # pylint: disable=missing-docstring
 
@@ -45,6 +45,7 @@ class VmSquash(cliapp.Application):
             description, epilog)
         self.rootdir = None
         self.filesystem = Filesystem()
+        self.remove_dirs = []
 
     def add_settings(self):
         self.settings.boolean(['verbose'], 'report what is going on')
@@ -53,33 +54,57 @@ class VmSquash(cliapp.Application):
         self.settings.string(['roottype'], 'specify file system type for /', default='ext4')
         self.settings.string(['tarball'], "tarball of the filesystem",
                              metavar='FILE')
+        self.settings.string(['directory'], 'unpacked rootfs directory.',
+                             metavar='DIRECTORY')
 
     def message(self, msg):
         logging.info(msg)
         if self.settings['verbose']:
             print msg
 
+    def cleanup_system(self):
+        # Clean up after any errors.
+
+        self.message('Cleaning up')
+
+        for dirname in self.remove_dirs:
+            shutil.rmtree(dirname)
+
     def process_args(self, args):
         """
-        Unpack a tarball of the disk's contents,
+        Optionally unpack a tarball of the disk's contents,
         shell out to runcmd since it more easily handles rootdir.
         Then run the vmdebootstrap Filesystem squash_rootfs.
         """
-        # unpacking tarballs containing device nodes needs root
-        if os.geteuid() != 0:
-            sys.exit("You need to have root privileges to run this script.")
-        self.filesystem.define_settings(self.settings)
-        rootdir = tempfile.mkdtemp()
-        logging.debug('mkdir %s', rootdir)
-        self.message('Unpacking tarball of disk contents')
-        self.filesystem.devices['rootdir'] = rootdir
-        runcmd(['tar', '-xf', self.settings['tarball'], '-C', rootdir])
-        # set rootdir in self.devices
-        self.devices = {
-            'rootdir': rootdir,
-        }
-        self.filesystem.squash_rootfs()
-        shutil.rmtree(rootdir)
+        if self.settings['directory'] and self.settings['tarball']:
+            raise cliapp.AppException(
+                'tarball and directory cannot be used together.')
+        if not self.settings['directory'] and not self.settings['tarball']:
+            raise cliapp.AppException(
+                'Specify either directory or a tarball.')
+        try:
+            self.filesystem.define_settings(self.settings)
+            if self.settings['tarball']:
+                # unpacking tarballs containing device nodes needs root
+                if os.geteuid() != 0:
+                    sys.exit("You need to have root privileges to unpack the tarball.")
+                rootdir = tempfile.mkdtemp()
+                self.remove_dirs.append(rootdir)
+                logging.debug('mkdir %s', rootdir)
+                self.message('Unpacking tarball of disk contents')
+                self.filesystem.devices['rootdir'] = rootdir
+                runcmd(['tar', '-xf', self.settings['tarball'], '-C', rootdir])
+            else:
+                self.message("Using %s directory" % self.settings['directory'])
+                self.filesystem.devices['rootdir'] = self.settings['directory']
+            self.filesystem.squash_rootfs()
+        except BaseException as e:
+            self.message('EEEK! Something bad happened...')
+            self.message(e)
+            self.cleanup_system()
+            raise
+        else:
+            self.cleanup_system()
 
 
 def main():