Modify rawdisk.write to allow the deployment to devices

This patch also modifies rawdisk.write to use the UPGRADE environment
variable to figure out when is doing an upgrade or a fresh deployment.
This change is important because os.path.isfile doesn't work with
devices.

1 files changed, 14 insertions, 6 deletions

diff --git a/rawdisk.write b/rawdisk.write
index 12db4398..e1a75fe0 100755
--- a/rawdisk.write
+++ b/rawdisk.write
@@ -43,17 +43,25 @@ class RawDiskWriteExtension(morphlib.writeexts.WriteExtension):
             raise cliapp.AppException('Wrong number of command line args')
         
         temp_root, location = args
-        if os.path.isfile(location):
+        upgrade = self.get_environment_boolean('UPGRADE')
+
+        if upgrade:
             self.upgrade_local_system(location, temp_root)
         else:
             try:
-                self.create_local_system(temp_root, location)
-                self.status(msg='Disk image has been created at %s' % location)
+                if not self.is_device(location):
+                    with self.created_disk_image(location):
+                        self.format_btrfs(location)
+                        self.create_system(temp_root, location)
+                    self.status(msg='Disk image has been created at %s' %
+                                     location)
+                else:
+                    self.format_btrfs(location)
+                    self.create_system(temp_root, location)
+                    self.status(msg='System deployed to %s' % location)
             except Exception:
-                self.status(msg='Failure to create disk image at %s' %
+                self.status(msg='Failure to deploy system to %s' %
                                  location)
-                if os.path.exists(location):
-                    os.remove(location)
                 raise
 
     def upgrade_local_system(self, raw_disk, temp_root):