reboot when done

1 files changed, 44 insertions, 10 deletions

diff --git a/loader.py b/loader.py
index 8d07e0a..c10b7bf 100755
--- a/loader.py
+++ b/loader.py
@@ -30,15 +30,27 @@ def read_stderr_realtime(proc, stream='stderr'):
 
 def check_environment(d):
     d.set_background_title("Checking your environment")
-    if subprocess.check_call(["ping", "8.8.8.8", "-c1", "-W5"]) != 0:
+
+    try:
+        subprocess.check_call(["ping", "8.8.8.8", "-c1", "-W5"])
+
+    except subprocess.CalledProcessError:
         d.set_background_title("Fatal: Not Connected To Network")
         d.msgbox("Please check your network connection and restart") 
         sys.exit(1)
-    if subprocess.check_call(["ping", "192.168.4.200", "-c1", "-W5"]) != 0:
+
+    try:
+        subprocess.check_call(["ping", "192.168.4.200", "-c1", "-W5"])
+
+    except subprocess.CalledProcessError:
         d.set_background_title("Fatal: Can not find depot")
         d.msgbox("Please check your network connection and restart") 
         sys.exit(1)
-    if subprocess.check_call(["mount","depot:/home/public/images","/mnt/images"]) !=0:
+        
+    try:
+        subprocess.check_call(["mount","depot:/home/public/images","/mnt/images"])
+ 
+    except subprocess.CalledProcessError:   
         d.set_background_title("Fatal: Can not mount images directory")
         d.msgbox("Please contact your system administrator") 
         sys.exit(1)
@@ -50,14 +62,18 @@ check_environment(d)
 d.set_background_title("Select file to load.....")
 
 code, selection = d.fselect('/mnt/images/',10,70)
+
 if code == d.OK :
-   
+    
+    disk_loaded=False
+       
     cmd = ['pv', '-n', selection ]
     try:
         output=open('/dev/sda','w')
     except:
-        print("Can't open target disk" )
-        pass
+        d.set_background_title("Fatal: Can not open target disk")
+        d.msgbox("Please check your system") 
+        sys.exit(1)
     try:
         proc = subprocess.Popen(
             cmd,
@@ -76,11 +92,29 @@ if code == d.OK :
                 break
             d.gauge_update(int(line.strip()))  
             if int(line.strip()) == 100:
-                d.gauge_update(100, "Finished!")       
+                d.gauge_update(100, "Finished!")
+                disk_loaded=True       
                 break
+
     except IOError:
-        print ('oops')
-        
-    exit_code = d.gauge_stop() 
+        d.set_background_title("Fatal: Could Not Complete Task")
+        d.msgbox("Sorry") 
+        sys.exit(1)
+
+    code = d.gauge_stop() 
 
+    if disk_loaded:
+        d.set_background_title("Success: Disk has been reimaged!")
+        code = d.pause("Please remove boot media, rebooting in 5 seconds", 10,70, 5)
+        if code=d.OK:
+            subprocess.check_call(["reboot"])
+        else:         
+            d.set_background_title("Reboot Cancelled")
+            d.msgbox("Have A Nice Day") 
+            sys.exit(2)
+else:
+     d.set_background_title("Cancelled")
+     d.msgbox("Have A Nice Day") 
+     sys.exit(2)
+