drakrun: Ensure OS_INFO.json exists before accessing it (#658)

drakrun/drakrun/draksetup.py

@@ -50,8 +50,19 @@
     start_dnsmasq,
     stop_dnsmasq,
 )
-from drakrun.storage import REGISTERED_BACKEND_NAMES, get_storage_backend
-from drakrun.util import RuntimeInfo, VmiOffsets, file_sha256, safe_delete
+from drakrun.storage import (
+    REGISTERED_BACKEND_NAMES,
+    StorageBackendBase,
+    get_storage_backend,
+)
+from drakrun.util import (
+    RuntimeInfo,
+    VmiGuidInfo,
+    VmiOffsets,
+    file_sha256,
+    safe_delete,
+    vmi_win_guid,
+)
 from drakrun.vm import (
     FIRST_CDROM_DRIVE,
     SECOND_CDROM_DRIVE,
@@ -697,6 +708,20 @@ def insert_cd(domain, drive, iso):
     subprocess.run(["xl", "cd-insert", domain, drive, iso], check=True)
 
 
+def build_os_info(
+    apiscout_profile_dir: str,
+    kernel_info: VmiGuidInfo,
+    storage_backend: StorageBackendBase,
+):
+    os_info = {
+        "os_name": kernel_info.version,
+        "os_timestamp": storage_backend.get_vm0_snapshot_time(),
+    }
+
+    with open(os.path.join(apiscout_profile_dir, "OS_INFO.json"), "w") as f:
+        f.write(json.dumps(os_info, indent=4, sort_keys=True))
+
+
 @click.command(help="Finalize sandbox installation")
 @click.option(
     "--report/--no-report",
@@ -737,23 +762,13 @@ def postinstall(report, generate_usermode):
         # If unattended install is enabled, we have an additional CD-ROM drive
         eject_cd("vm-0", SECOND_CDROM_DRIVE)
 
-    output = subprocess.check_output(
-        ["vmi-win-guid", "name", "vm-0"], timeout=30
-    ).decode("utf-8")
+    kernel_info = vmi_win_guid("vm-0")
 
-    try:
-        version = re.search(r"Version: (.*)", output).group(1)
-        pdb = re.search(r"PDB GUID: ([0-9a-f]+)", output).group(1)
-        fn = re.search(r"Kernel filename: ([a-z]+\.[a-z]+)", output).group(1)
-    except AttributeError:
-        logging.error("Failed to obtain kernel PDB GUID/Kernel filename.")
-        return
-
-    logging.info("Determined PDB GUID: {}".format(pdb))
-    logging.info("Determined kernel filename: {}".format(fn))
+    logging.info(f"Determined PDB GUID: {kernel_info.guid}")
+    logging.info(f"Determined kernel filename: {kernel_info.filename}")
 
     logging.info("Fetching PDB file...")
-    dest = fetch_pdb(fn, pdb, destdir=PROFILE_DIR)
+    dest = fetch_pdb(kernel_info.filename, kernel_info.guid, destdir=PROFILE_DIR)
 
     logging.info("Generating profile out of PDB file...")
     profile = make_pdb_profile(dest)
@@ -788,19 +803,15 @@ def postinstall(report, generate_usermode):
     if report:
         send_usage_report(
             {
-                "kernel": {"guid": pdb, "filename": fn, "version": version},
+                "kernel": {
+                    "guid": kernel_info.guid,
+                    "filename": kernel_info.filename,
+                    "version": kernel_info.version,
+                },
                 "install_iso": {"sha256": install_info.iso_sha256},
             }
         )
 
-    os_info = {
-        "os_name": version,
-        "os_timestamp": storage_backend.get_vm0_snapshot_time(),
-    }
-
-    with open(os.path.join(APISCOUT_PROFILE_DIR, "OS_INFO.json"), "w") as f:
-        f.write(json.dumps(os_info, indent=4, sort_keys=True))
-
     if generate_usermode:
         # Restore a VM and create usermode profiles
         create_missing_profiles()
@@ -859,6 +870,8 @@ def create_missing_profiles():
             except Exception:
                 log.exception("Unexpected exception from create_rekall_profile!")
 
+    build_os_info(APISCOUT_PROFILE_DIR, vmi_win_guid(vm.vm_name), backend)
+
     dll_basename_list = [dll.dest for dll in dll_file_list]
     static_apiscout_profile = build_static_apiscout_profile(
         APISCOUT_PROFILE_DIR, dll_basename_list

drakrun/drakrun/util.py

@@ -223,3 +223,29 @@ def file_sha256(filename, blocksize=65536) -> str:
         for block in iter(lambda: f.read(blocksize), b""):
             file_hash.update(block)
     return file_hash.hexdigest()
+
+
+@dataclass
+class VmiGuidInfo:
+    version: str
+    guid: str
+    filename: str
+
+
+def vmi_win_guid(vm_name: str) -> VmiGuidInfo:
+    result = subprocess.run(
+        ["vmi-win-guid", "name", vm_name],
+        timeout=30,
+        capture_output=True,
+    )
+
+    output = result.stdout.decode()
+
+    version = re.search(r"Version: (.*)", output)
+    pdb_guid = re.search(r"PDB GUID: ([0-9a-f]+)", output)
+    kernel_filename = re.search(r"Kernel filename: ([a-z]+\.[a-z]+)", output)
+
+    if version is None or pdb_guid is None or kernel_filename is None:
+        raise RuntimeError("Invalid vmi-win-guid output")
+
+    return VmiGuidInfo(version.group(1), pdb_guid.group(1), kernel_filename.group(1))