Add ApiScout profile to do_export_full and do_import_full

drakrun/drakrun/apiscout.py

@@ -55,8 +55,6 @@ def make_static_apiscout_profile_for_dll(filepath: str) -> Dict[str, Any]:
             pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"],
         ]
     )
-    if not hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
-        raise RuntimeError(f"DIRECTORY_ENTRY_EXPORT not found in '{filepath}'")
 
     dll_entry = {}
     dll_entry["base_address"] = pe.OPTIONAL_HEADER.ImageBase
@@ -65,6 +63,10 @@ def make_static_apiscout_profile_for_dll(filepath: str) -> Dict[str, Any]:
     dll_entry["filepath"] = filepath
     dll_entry["aslr_offset"] = 0
     dll_entry["exports"] = []
+    if not hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
+        if pe.is_driver():
+            return dll_entry
+        raise RuntimeError(f"DIRECTORY_ENTRY_EXPORT not found in '{filepath}'")
     for exp in sorted(pe.DIRECTORY_ENTRY_EXPORT.symbols, key=attrgetter("address")):
         export_info = {}
 
@@ -96,6 +98,11 @@ def build_static_apiscout_profile(
 
     for dll_basename in dll_basename_list:
         filepath = Path(apiscout_profile_dir) / f"{dll_basename}.json"
+        if not filepath.is_file():
+            log.warning(
+                f"'{filepath}' not found. Is there a problem with profiles generation?"
+            )
+            continue
         with open(filepath) as f:
             dll_profile = json.load(f)
         dlls_profiles[build_apiscout_dll_key(dll_profile)] = dll_profile

drakrun/drakrun/draksetup.py

@@ -22,7 +22,10 @@
 from requests import RequestException
 from tqdm import tqdm
 
-from drakrun.apiscout import make_static_apiscout_profile_for_dll
+from drakrun.apiscout import (
+    build_static_apiscout_profile,
+    make_static_apiscout_profile_for_dll,
+)
 from drakrun.config import (
     APISCOUT_PROFILE_DIR,
     ETC_DIR,
@@ -779,10 +782,6 @@ def postinstall(report, generate_usermode):
     logging.info("Snapshotting persistent memory...")
     storage_backend.snapshot_vm0_volume()
 
-    if generate_usermode:
-        # Restore a VM and create usermode profiles
-        create_missing_profiles()
-
     if report:
         send_usage_report(
             {
@@ -799,6 +798,10 @@ def postinstall(report, generate_usermode):
     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()
+
     logging.info("All right, drakrun setup is done.")
     logging.info("First instance of drakrun will be enabled automatically...")
     subprocess.check_output("systemctl enable drakrun@1", shell=True)
@@ -808,8 +811,10 @@ def postinstall(report, generate_usermode):
     logging.info("  # draksetup scale <number of instances>")
 
 
-def profile_exists(profile: DLL) -> bool:
-    return (Path(PROFILE_DIR) / f"{profile.dest}.json").is_file()
+def profiles_exist(profile_name: str) -> bool:
+    return (Path(PROFILE_DIR) / f"{profile_name}.json").is_file() and (
+        Path(APISCOUT_PROFILE_DIR) / f"{profile_name}.json"
+    ).is_file()
 
 
 def create_missing_profiles():
@@ -841,16 +846,23 @@ def create_missing_profiles():
     # Ensure that all declared usermode profiles exist
     # This is important when upgrade defines new entries in dll_file_list and compulsory_dll_file_list
     for profile in compulsory_dll_file_list:
-        if not profile_exists(profile):
+        if not profiles_exist(profile.dest):
             create_rekall_profile(injector, profile, True)
 
     for profile in dll_file_list:
-        if not profile_exists(profile):
+        if not profiles_exist(profile.dest):
             try:
                 create_rekall_profile(injector, profile)
             except Exception:
                 log.exception("Unexpected exception from create_rekall_profile!")
 
+    dll_basename_list = [dll.dest for dll in dll_file_list]
+    static_apiscout_profile = build_static_apiscout_profile(
+        APISCOUT_PROFILE_DIR, dll_basename_list
+    )
+    with open(Path(APISCOUT_PROFILE_DIR) / "static_apiscout_profile.json", "w") as f:
+        json.dump(static_apiscout_profile, f)
+
     vm.destroy()
     delete_vm_network(
         vm_id=1, net_enable=False, out_interface=out_interface, dns_server=dns_server
@@ -1277,6 +1289,15 @@ def do_export_full(mc, bucket, name):
             bucket, f"{name}/profiles/{file}", os.path.join(PROFILE_DIR, file)
         )
 
+    # Upload ApiScout profile
+    for file in os.listdir(APISCOUT_PROFILE_DIR):
+        logging.info("Uploading file %s", file)
+        mc.fput_object(
+            bucket,
+            f"{name}/apiscout_profile/{file}",
+            os.path.join(APISCOUT_PROFILE_DIR, file),
+        )
+
 
 def do_import_full(mc, name, bucket, zpool):
     """Perform full snapshot import, symmetric to do_export_full"""
@@ -1291,14 +1312,22 @@ def do_import_full(mc, name, bucket, zpool):
                 ["zcat", compressed_snapshot.name], stdout=snapshot, check=True
             )
 
-    profile_prefix = f"{name}/profiles/"
-    for object in mc.list_objects(bucket, prefix=profile_prefix):
-        # Strip profile prefix
-        profile_name = object.object_name[len(profile_prefix) :]
+    profiles_prefix = f"{name}/profiles/"
+    for object in mc.list_objects(bucket, prefix=profiles_prefix):
+        # Strip profiles prefix
+        profile_name = object.object_name[len(profiles_prefix) :]
         mc.fget_object(
             bucket, object.object_name, os.path.join(PROFILE_DIR, profile_name)
         )
 
+    apiscout_profile_prefix = f"{name}/apiscout_profile/"
+    for object in mc.list_objects(bucket, prefix=apiscout_profile_prefix):
+        # Strip apiscout profile prefix
+        filename = object.object_name[len(apiscout_profile_prefix) :]
+        mc.fget_object(
+            bucket, object.object_name, os.path.join(APISCOUT_PROFILE_DIR, filename)
+        )
+
 
 @click.group()
 def main():

drakrun/drakrun/main.py

@@ -27,7 +27,6 @@
 from karton.core import Config, Karton, LocalResource, Resource, Task
 
 import drakrun.sample_startup as sample_startup
-from drakrun.apiscout import build_static_apiscout_profile
 from drakrun.config import (
     APISCOUT_PROFILE_DIR,
     ETC_DIR,
@@ -387,16 +386,6 @@ def build_profile_payload(self) -> Dict[str, LocalResource]:
 
             return Resource.from_directory(name="profiles", directory_path=tmp_dir)
 
-    def build_static_apiscout_profile_payload(self) -> Dict[str, LocalResource]:
-        dll_basename_list = [dll.dest for dll in dll_file_list]
-        static_apiscout_profile = build_static_apiscout_profile(
-            APISCOUT_PROFILE_DIR, dll_basename_list
-        )
-        return LocalResource(
-            name="static_apiscout_profile.json",
-            content=json.dumps(static_apiscout_profile, indent=4, sort_keys=True),
-        )
-
     def send_raw_analysis(self, sample, outdir, metadata, dumps_metadata, quality):
         """
         Offload drakrun-prod by sending raw analysis output to be processed by
@@ -429,7 +418,10 @@ def send_raw_analysis(self, sample, outdir, metadata, dumps_metadata, quality):
             self.log.info("Uploading static ApiScout profile...")
             task.add_payload(
                 "static_apiscout_profile.json",
-                self.build_static_apiscout_profile_payload(),
+                LocalResource(
+                    name="static_apiscout_profile.json",
+                    path=Path(APISCOUT_PROFILE_DIR) / "static_apiscout_profile.json",
+                ),
             )
 
         self.log.info("Uploading artifacts...")