Refactored cleanup

1 files changed, 12 insertions, 25 deletions

diff --git a/install b/install
index 30ce0c8..3157324 100755
--- a/install
+++ b/install
@@ -659,23 +659,6 @@ def umount_in_bootenv():
     for fs in fslist:
         call(['umount', rootdir + fs], stdout=PIPE, stderr=PIPE)
 
-def destroy_bootenv():
-    global bootenv
-    global rootdir
-    if not bootenv:
-        return
-    umount_in_bootenv()
-    progress = ProgressMessage(screen, title='Destroying boot environment',
-            text='Destroying "{}", please wait ...'.format(bootenv))
-    call(['umount', rootdir], stdout=PIPE, stderr=PIPE)
-    call(['zfs', 'destroy', '-r',  bootenv], stdout=PIPE, stderr=PIPE)
-    bootenv = None
-    try:
-        os.rmdir(rootdir)
-        rootdir = None
-    except:
-        pass
-
 def in_bootenv(cmd):
     chroot = ['chroot', rootdir]
     chroot += cmd
@@ -752,15 +735,19 @@ def configure_bootenv():
     create_bootarchive()
 
 
-def cleanup():
-    '''Cleanup after successful installation'''
-    progress = ProgressMessage(screen, title='Cleaning up', width=60)
-    progress.text = 'Unmounting boot environment {} ...'.format(bootenv)
+def cleanup(destroy_bootenv=False):
+    progress = ProgressMessage(screen, title='Cleaning up', width=50)
     umount_in_bootenv()
     call(['umount', rootdir], stdout=PIPE, stderr=PIPE)
-    call(['zfs', 'destroy', '-r',  bootenv + '@empty'], stdout=PIPE, stderr=PIPE)
-    call(['zfs', 'set', 'canmount=noauto', bootenv], stdout=PIPE, stderr=PIPE)
-    call(['zfs', 'set', 'mountpoint=/', bootenv], stdout=PIPE, stderr=PIPE)
+    if destroy_bootenv:
+        progress.text='Destroying "{}", please wait ...'.format(bootenv)
+        call(['zfs', 'destroy', '-r',  bootenv], stdout=PIPE, stderr=PIPE)
+    else:
+        progress.text = 'Adjusting boot environment "{}" ...'.format(bootenv)
+        call(['zfs', 'destroy', '-r',  bootenv + '@empty'], stdout=PIPE, stderr=PIPE)
+        call(['zfs', 'set', 'canmount=noauto', bootenv], stdout=PIPE, stderr=PIPE)
+        call(['zfs', 'set', 'mountpoint=/', bootenv], stdout=PIPE, stderr=PIPE)
+
     try:
         os.rmdir(rootdir)
     except:
@@ -789,7 +776,7 @@ try:
     goodbye()
 
 except Abort as e:
-    destroy_bootenv()
+    cleanup(destroy_bootenv=True)
 except NoDisks as e:
     print (e)
 finally: