drakrun: Dump raw guest VM memory feature (#621)

drakrun/drakrun/config.dist.ini

@@ -23,6 +23,10 @@ access_key=
 secret_key=
 
 [drakrun]
+; if set, drakrun dumps raw VM guest memory after finished analysis with drakvuf
+; compressed dumps are stored as "post_sample.raw_memdump.gz" in analysis folder
+raw_memory_dump=0
+
 ; whether guest VMs should have access to the Internet or no
 net_enable=0
 

drakrun/drakrun/draksetup.py

@@ -52,7 +52,7 @@
     VirtualMachine,
 )
 from drakrun.apiscout import make_static_apiscout_profile_for_dll
-from drakrun.util import RuntimeInfo, VmiOffsets, safe_delete
+from drakrun.util import RuntimeInfo, VmiOffsets, safe_delete, file_sha256
 from tqdm import tqdm
 from pathlib import Path, PureWindowsPath
 import traceback
@@ -1004,6 +1004,75 @@ def get_minio_client(config):
     )
 
 
+@click.group(help="Manage VM raw memory pre-sample dumps")
+def memdump():
+    pass
+
+
+@memdump.command(name="export", help="Upload pre-sample raw memory dump to MinIO.")
+@click.option("--instance", required=True, type=int, help="Instance ID of restored VM")
+@click.option(
+    "--bucket",
+    default="presample-memdumps",
+    help="MinIO bucket to store the compressed raw image",
+)
+def memdump_export(bucket, instance):
+    install_info = InstallInfo.try_load()
+    if install_info is None:
+        logging.error(
+            "Missing installation info. Did you forget to set up the sandbox?"
+        )
+        return
+
+    backend = get_storage_backend(install_info)
+    vm = VirtualMachine(backend, instance)
+    if vm.is_running:
+        logging.exception(f"vm-{instance} is running")
+        return
+
+    logging.info("Calculating snapshot hash...")
+    snapshot_sha256 = file_sha256(os.path.join(VOLUME_DIR, "snapshot.sav"))
+    name = f"{snapshot_sha256}_pre_sample.raw_memdump.gz"
+
+    mc = get_minio_client(conf)
+
+    if not mc.bucket_exists(bucket):
+        logging.error("Bucket %s doesn't exist", bucket)
+        return
+
+    try:
+        mc.stat_object(bucket, name)
+        logging.info("This file already exists in specified bucket")
+        return
+    except NoSuchKey:
+        pass
+    except Exception:
+        logging.exception("Failed to check if object exists on minio")
+
+    logging.info("Restoring VM and performing memory dump")
+
+    try:
+        vm.restore(pause=True)
+    except subprocess.CalledProcessError:
+        logging.exception(f"Failed to restore VM {vm.vm_name}")
+        with open(f"/var/log/xen/qemu-dm-{vm.vm_name}.log", "rb") as f:
+            logging.error(f.read())
+    logging.info("VM restored")
+
+    with tempfile.NamedTemporaryFile() as compressed_memdump:
+        vm.memory_dump(compressed_memdump.name)
+
+        logging.info(f"Uploading {name} to {bucket}")
+        mc.fput_object(bucket, name, compressed_memdump.name)
+
+    try:
+        vm.destroy()
+    except Exception:
+        logging.exception("Failed to destroy VM")
+
+    logging.info("Done")
+
+
 @click.group(help="Manage VM snapshots")
 def snapshot():
     pass
@@ -1247,6 +1316,7 @@ def main():
 main.add_command(mount)
 main.add_command(scale)
 main.add_command(snapshot)
+main.add_command(memdump)
 main.add_command(cleanup)
 
 

drakrun/drakrun/main.py

@@ -28,7 +28,13 @@
 
 from drakrun.version import __version__ as DRAKRUN_VERSION
 from drakrun.drakpdb import dll_file_list
-from drakrun.config import InstallInfo, ETC_DIR, PROFILE_DIR, APISCOUT_PROFILE_DIR
+from drakrun.config import (
+    InstallInfo,
+    ETC_DIR,
+    PROFILE_DIR,
+    APISCOUT_PROFILE_DIR,
+    VOLUME_DIR,
+)
 from drakrun.storage import get_storage_backend
 from drakrun.networking import start_tcpdump_collector, start_dnsmasq, setup_vm_network
 from drakrun.util import (
@@ -37,6 +43,7 @@
     get_xen_commandline,
     graceful_exit,
     RuntimeInfo,
+    file_sha256,
 )
 from drakrun.vm import generate_vm_conf, VirtualMachine
 from drakrun.injector import Injector
@@ -248,6 +255,7 @@ def vm_name(self) -> str:
         return f"vm-{self.instance_id}"
 
     def init_drakrun(self):
+
         generate_vm_conf(self.install_info, self.instance_id)
 
         if not self.backend.minio.bucket_exists("drakrun"):
@@ -258,6 +266,9 @@ def init_drakrun(self):
 
         setup_vm_network(self.instance_id, self.net_enable, out_interface, dns_server)
 
+        self.log.info("Caculating snapshot hash...")
+        self.snapshot_sha256 = file_sha256(os.path.join(VOLUME_DIR, "snapshot.sav"))
+
     def _karton_safe_get_headers(self, task, key, fallback):
         ret = task.headers.get(key, fallback)
         # intentional workaround due to a bug in karton
@@ -570,6 +581,9 @@ def analyze_sample(self, sample_path, workdir, outdir, start_command, timeout):
 
         dns_server = self.config.config["drakrun"].get("dns_server", "8.8.8.8")
         drakmon_log_fp = os.path.join(outdir, "drakmon.log")
+        raw_memory_dump = self.config.config["drakrun"].getboolean(
+            "raw_memory_dump", fallback=False
+        )
 
         with self.run_vm() as vm, graceful_exit(
             start_dnsmasq(self.instance_id, dns_server)
@@ -642,6 +656,9 @@ def analyze_sample(self, sample_path, workdir, outdir, start_command, timeout):
                 self.log.exception("DRAKVUF timeout expired")
                 raise e
 
+            if raw_memory_dump:
+                vm.memory_dump(os.path.join(outdir, "post_sample.raw_memdump.gz"))
+
         return analysis_info
 
     @with_logs("drakrun.log")
@@ -654,6 +671,7 @@ def process(self, task: Task):
         self.log.info(f"Running on: {socket.gethostname()}")
         self.log.info(f"Sample SHA256: {sha256sum}")
         self.log.info(f"Analysis UID: {self.analysis_uid}")
+        self.log.info(f"Snapshot SHA256: {self.snapshot_sha256}")
 
         # Timeout sanity check
         timeout = task.payload.get("timeout") or self.default_timeout
@@ -717,6 +735,7 @@ def process(self, task: Task):
 
         metadata = {
             "sample_sha256": sha256sum,
+            "snapshot_sha256": self.snapshot_sha256,
             "magic_output": magic_output,
             "time_started": int(time.time()),
         }

drakrun/drakrun/util.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass, field
 from typing import IO, AnyStr
 import traceback
+import hashlib
 
 from dataclasses_json import config, dataclass_json
 
@@ -214,3 +215,11 @@ def graceful_exit(proc: subprocess.Popen):
             proc.kill()
             proc.wait()
             log.error("Process was forceully killed")
+
+
+def file_sha256(filename, blocksize=65536) -> str:
+    file_hash = hashlib.sha256()
+    with open(filename, "rb") as f:
+        for block in iter(lambda: f.read(blocksize), b""):
+            file_hash.update(block)
+    return file_hash.hexdigest()

drakrun/drakrun/vm.py

@@ -2,6 +2,7 @@
 import re
 import logging
 import subprocess
+import tempfile
 from pathlib import Path
 
 from drakrun.storage import get_storage_backend, StorageBackendBase
@@ -174,3 +175,30 @@ def destroy(self, **kwargs):
                 f"Failed to destroy VM {self.vm_name}",
                 **kwargs,
             )
+
+    def memory_dump(self, compressed_filepath):
+        """Dump raw memory from running vm using vmi-dump-memory and compress it with gzip
+        :raises: subprocess.CalledProcessError
+        """
+
+        with tempfile.NamedTemporaryFile() as raw_memdump, open(
+            compressed_filepath, "wb"
+        ) as compressed_file:
+
+            log.info(f"Dumping raw memory from {self.vm_name} guest...")
+            try:
+                subprocess.run(
+                    ["vmi-dump-memory", self.vm_name, raw_memdump.name], check=True
+                )
+            except subprocess.CalledProcessError as e:
+                log.error(f"Dumping raw memory from {self.vm_name} failed.")
+                raise e
+
+            log.info(f"Compressing {self.vm_name} guest memory dump...")
+            try:
+                subprocess.run(
+                    ["gzip", "-c", raw_memdump.name], check=True, stdout=compressed_file
+                )
+            except subprocess.CalledProcessError as e:
+                log.error(f"Compressing raw memory from {self.vm_name} failed.")
+                raise e