Recover from fail in vm-1 stage

CERT-Polska · Sep 25, 2023 · 8c79e93 · 8c79e93
1 parent 3073e36
commit 8c79e93
Showing 1 changed file with 34 additions and 28 deletions.
diff --git a/draklib/draklib/draksh/postinstall.py b/draklib/draklib/draksh/postinstall.py
@@ -14,7 +14,7 @@
     vmi_win_guid,
 )
 from ..drakvuf.vm import DrakvufVM
-from ..machinery.vm import FIRST_CDROM_DRIVE, SECOND_CDROM_DRIVE
+from ..machinery.vm import VirtualMachine, FIRST_CDROM_DRIVE, SECOND_CDROM_DRIVE
 from ..util import ensure_delete
 from .util import check_root
 
@@ -37,40 +37,46 @@ def postinstall(profile_name):
     profile = Profile.load(profile_name)
 
     vm0 = DrakvufVM(profile, 0)
+    vm1 = DrakvufVM(profile, 1)
 
-    if vm0.vm.is_running is False:
-        raise click.ClickException("vm-0 is not running")
+    if vm1.vm.is_running is True:
+        # If vm1 is running: probably we failed to make a DLL profile
+        # Let's revert it and use already made snapshot
+        vm1.vm.destroy()
+    else:
+        if vm0.vm.is_running is False:
+            raise click.ClickException("vm-0 is not running")
 
-    logging.info("Cleaning up leftovers (if any)")
-    for path in profile.vm_profile_dir.glob("*"):
-        ensure_delete(path)
+        logging.info("Cleaning up leftovers (if any)")
+        for path in profile.vm_profile_dir.glob("*"):
+            ensure_delete(path)
 
-    logging.info("Ejecting installation CDs")
-    vm0.vm.eject_cd(FIRST_CDROM_DRIVE)
-    if profile.install_info.enable_unattended:
-        # If unattended installation is enabled, we have an additional CD-ROM drive
-        # TODO
-        vm0.vm.eject_cd(SECOND_CDROM_DRIVE)
+        logging.info("Ejecting installation CDs")
+        vm0.vm.eject_cd(FIRST_CDROM_DRIVE)
+        if profile.install_info.enable_unattended:
+            # If unattended installation is enabled, we have an additional CD-ROM drive
+            # TODO
+            vm0.vm.eject_cd(SECOND_CDROM_DRIVE)
 
-    win_guid_info = vm0.get_win_guid()
+        win_guid_info = vm0.get_win_guid()
 
-    logging.info(f"Determined Windows version: {win_guid_info.version}")
-    logging.info(f"Determined PDB GUID: {win_guid_info.guid}")
-    logging.info(f"Determined kernel filename: {win_guid_info.filename}")
+        logging.info(f"Determined Windows version: {win_guid_info.version}")
+        logging.info(f"Determined PDB GUID: {win_guid_info.guid}")
+        logging.info(f"Determined kernel filename: {win_guid_info.filename}")
 
-    vm0.create_kernel_profile(win_guid_info)
-    vm0.create_runtime_info()
+        vm0.create_kernel_profile(win_guid_info)
+        vm0.create_runtime_info()
 
-    logging.info("Saving VM snapshot...")
-    # Create vm-0 snapshot, and destroy it
-    # WARNING: qcow2 snapshot method is a noop. fresh images are created on the fly
-    # so we can't keep the vm-0 running
-    vm0.save(destroy_after=True)
-    logging.info("Snapshot was saved succesfully.")
+        logging.info("Saving VM snapshot...")
+        # Create vm-0 snapshot, and destroy it
+        # WARNING: qcow2 snapshot method is a noop. fresh images are created on the fly
+        # so we can't keep the vm-0 running
+        vm0.save(destroy_after=True)
+        logging.info("Snapshot was saved succesfully.")
 
-    # Memory state is frozen, we can't do any writes to persistent storage
-    logging.info("Snapshotting persistent memory...")
-    vm0.vm.storage.snapshot_vm0_volume()
+        # Memory state is frozen, we can't do any writes to persistent storage
+        logging.info("Snapshotting persistent memory...")
+        vm0.vm.storage.snapshot_vm0_volume()
 
     # Restore a VM and create DLL profiles
     vm1 = DrakvufVM(profile, 1)
@@ -87,7 +93,7 @@ def postinstall(profile_name):
     for dllspec in optional_dlls:
         try:
             vm1.make_dll_profile(dllspec)
-        except InjectorError:
+        except Exception:
             log.exception("Failed to profile optional DLL")
             failed_dlls.append(dllspec)