From fc702e29b8f2d32c159ebf4ddaa33238607b934f Mon Sep 17 00:00:00 2001
From: dbarzin <dbarzin@gmail.com>
Date: Tue, 7 Mar 2023 11:25:00 +0100
Subject: [PATCH] fix invalid labels

---
 pandora-box.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/pandora-box.py b/pandora-box.py
index fa09885..02ccc40 100755
--- a/pandora-box.py
+++ b/pandora-box.py
@@ -728,7 +728,7 @@ def clean():
             log('PRESS KEY TO CLEAN', flush=True)
             screen.getch()
 
-        # TODO: check key is still here
+        # TODO: check device is still present
 
         # Remove infected files
         files_removed = 0
@@ -748,20 +748,26 @@ def clean():
                     f'not_removed="{file}, '
                     f'error="{str(ex)}"', exc_info=True)
                 has_error = True
-        if not has_curses:
-            display_image("OK")
-        else:
-            if not has_error:
+
+        umount_device()
+
+        if not has_error:
+            if has_curses:
                 log('Device cleaned !', flush=True)
             else:
+                display_image("OK")
+        else:
+            if has_curses:
                 log('Device not cleaned !', flush=True)
+            else:
+                display_image("ERROR")
+                wait_mouse_click()
         logging.info(
             f'hostname="{hostname}", '
             f'cleaned="{files_removed}/{len(infected_files)}"')
     else:
         if not has_curses:
             display_image("OK")
-    umount_device()
     return "WAIT"