Merge branch 'master' into feature/s3-storage

drakrun/analyzer/postprocessing/indexer.py

@@ -3,6 +3,7 @@ import pathlib
 from collections import defaultdict
 from typing import List, Optional
 
+import msgpack
 import orjson
 
 from .process_tree import ProcessTree
@@ -72,6 +73,58 @@ def index_log_file(
     return dict(index)
 
 
+def build_log_index(analysis_dir: pathlib.Path, process_tree: ProcessTree) -> bytes:
+    log_index = [dict() for _ in range(len(process_tree.processes))]
+    for log_file_path in analysis_dir.glob("*.log"):
+        plugin_name = log_file_path.stem
+        if plugin_name == "drakrun":
+            continue
+        index = index_log_file(
+            log_file_path,
+            process_tree,
+            key="Method",
+        )
+        for proc_seqid in index.keys():
+            log_index[proc_seqid][plugin_name] = index[proc_seqid]
+
+    serialized_indexes = []
+    serialized_indexes_current = 0
+    index_toc = [dict() for _ in range(len(process_tree.processes))]
+    for proc_seqid, proc_log_index in enumerate(log_index):
+        for plugin_name, plugin_index in proc_log_index.items():
+            serialized_index = msgpack.packb(plugin_index)
+            index_toc[proc_seqid][plugin_name] = serialized_indexes_current
+            serialized_indexes.append(serialized_index)
+            serialized_indexes_current += len(serialized_index)
+    serialized_index_toc = msgpack.packb(index_toc)
+    return serialized_index_toc + b"".join(serialized_indexes)
+
+
+def get_plugin_names_for_process(index_file: pathlib.Path, seqid: int):
+    with index_file.open("rb") as f:
+        unpacker = msgpack.Unpacker(f)
+        index_toc = unpacker.unpack()
+        if seqid >= len(index_toc):
+            return []
+        return list(index_toc[seqid].keys())
+
+
+def get_log_index_for_process(index_file: pathlib.Path, seqid: int, plugin_name: str):
+    with index_file.open("rb") as f:
+        unpacker = msgpack.Unpacker(f)
+        index_toc = unpacker.unpack()
+        log_pos = unpacker.tell()
+        if seqid >= len(index_toc):
+            return None
+        if plugin_name not in index_toc[seqid]:
+            return None
+        start = index_toc[seqid][plugin_name]
+        # Seek to the specific point of file
+        f.seek(log_pos + start)
+        unpacker = msgpack.Unpacker(f)
+        return unpacker.unpack()
+
+
 def scattered_read_file(
     log_file: pathlib.Path,
     offsets: List[List[int]],

drakrun/analyzer/postprocessing/plugins/__init__.py

@@ -50,6 +50,6 @@ POSTPROCESS_PLUGINS = [
     PostprocessPlugin(function=crop_dumps, requires=[DUMPS_DIR], generates=[DUMPS_ZIP]),
     PostprocessPlugin(function=compress_ipt, requires=[IPT_DIR], generates=[IPT_ZIP]),
     PostprocessPlugin(
-        function=index_logs, requires=["procmon.log"], generates=["index"]
+        function=index_logs, requires=["process_tree.json"], generates=["log_index"]
     ),
 ]

drakrun/analyzer/postprocessing/plugins/index_logs.py

@@ -1,31 +1,13 @@
 import json
 import pathlib
 
-from ..indexer import index_log_file
-from ..process_tree import tree_from_log
+from ..indexer import build_log_index
+from ..process_tree import tree_from_dict
 
 
 def index_logs(analysis_dir: pathlib.Path) -> None:
-    index_dir = analysis_dir / "index"
-    index_dir.mkdir(exist_ok=True)
-
-    # TODO: I need ProcessTree object but right now
-    #       I can't easily load it from process_tree.json
-
-    procmon_log_path = analysis_dir / "procmon.log"
-    with procmon_log_path.open("r") as procmon_log:
-        process_tree = tree_from_log(procmon_log)
-
-    for log_file_path in analysis_dir.glob("*.log"):
-        plugin_name = log_file_path.stem
-        if plugin_name == "drakrun":
-            continue
-        index = index_log_file(
-            log_file_path,
-            process_tree,
-            key="Method",
-        )
-
-        for seqid in index.keys():
-            index_path = index_dir / f"{plugin_name}.{seqid}.json"
-            index_path.write_text(json.dumps(index[seqid]))
+    process_tree_dict = json.loads((analysis_dir / "process_tree.json").read_text())
+    process_tree = tree_from_dict(process_tree_dict)
+    index = build_log_index(analysis_dir, process_tree)
+    index_path = analysis_dir / "log_index"
+    index_path.write_bytes(index)

drakrun/analyzer/worker.py

@@ -126,6 +126,7 @@ def enqueue_analysis(
     file_metadata: FileMetadata,
     options: AnalysisOptions,
     connection: Redis,
+    result_ttl: int,
 ) -> Job:
     queue = Queue(name=ANALYSIS_QUEUE_NAME, connection=connection)
     if options.sample_path is None:
@@ -141,5 +142,5 @@ def enqueue_analysis(
             "file": file_metadata.to_dict(),
         },
         job_timeout=options.timeout + options.job_timeout_leeway,
-        result_ttl=-1,
+        result_ttl=result_ttl,
     )

drakrun/lib/config.py

@@ -37,6 +37,7 @@ class DrakrunConfigSection(BaseModel):
     """Give extra 10 minutes as a timeout for whole analysis process
     including VM restore, post-restore, drakvuf hard timeout and
     postprocessing."""
+    results_ttl: int = -1
     apimon_hooks_path: Optional[pathlib.Path] = None
     syscall_hooks_path: Optional[pathlib.Path] = None
     extra_drakvuf_args: Optional[Dict[str, Any]] = None

drakrun/web/analysis.py

@@ -1,6 +1,10 @@
 import json
 import pathlib
 
+from drakrun.analyzer.postprocessing.indexer import (
+    get_log_index_for_process,
+    get_plugin_names_for_process,
+)
 from drakrun.analyzer.postprocessing.process_tree import tree_from_dict
 from drakrun.lib.paths import ANALYSES_DIR, DUMPS_ZIP
 
@@ -34,12 +38,12 @@ class AnalysisStorage:
         """Download DRAKVUF log"""
         return self._check_path(self.analysis_dir / f"{log_type}.log")
 
-    def get_log_index(self, log_type):
+    def get_log_index(self):
         """
         Download log index, useful for quickly accessing n-th
         log line
         """
-        return self._check_path(self.analysis_dir / "index" / log_type)
+        return self.analysis_dir / "log_index"
 
     def get_pcap_dump(self):
         """Download dump.pcap file."""
@@ -75,14 +79,16 @@ class AnalysisStorage:
             return None
         process_tree = tree_from_dict(json.loads(path.read_text()))
         process = process_tree.processes[which]
-        log_index = self.analysis_dir / "index"
+        log_index = self.analysis_dir / "log_index"
+        logs = {}
         if log_index.exists():
-            logs = {}
-            for index_path in log_index.glob(f"*.{which}.json"):
-                log_name = index_path.name.split(".")[0]
-                logs[log_name] = json.loads(index_path.read_text())["values"]
-        else:
-            logs = {}
+            plugin_names = get_plugin_names_for_process(log_index, which)
+            for plugin_name in plugin_names:
+                process_log_index = get_log_index_for_process(
+                    log_index, which, plugin_name
+                )
+                if process_log_index:
+                    logs[plugin_name] = process_log_index["values"]
         return {
             "process": process.as_dict(),
             "logs": logs,

drakrun/web/api.py

@@ -6,7 +6,6 @@ from tempfile import NamedTemporaryFile
 from zipfile import ZIP_DEFLATED, ZipFile
 
 import magic
-import orjson
 from flask import Response, jsonify, request, send_file
 from flask_openapi3 import APIBlueprint
 from rq.exceptions import NoSuchJobError
@@ -14,7 +13,10 @@ from rq.job import Job, JobStatus
 
 from drakrun.analyzer.analysis_options import AnalysisOptions
 from drakrun.analyzer.file_metadata import FileMetadata
-from drakrun.analyzer.postprocessing.indexer import scattered_read_file
+from drakrun.analyzer.postprocessing.indexer import (
+    get_log_index_for_process,
+    scattered_read_file,
+)
 from drakrun.analyzer.worker import (
     analysis_job_to_status_dict,
     enqueue_analysis,
@@ -89,6 +91,7 @@ def upload_sample(form: UploadFileForm):
             file_metadata=file_metadata,
             options=analysis_options,
             connection=redis,
+            result_ttl=config.drakrun.result_ttl,
         )
         add_analysis_to_recent(connection=redis, analysis_id=job_id)
         return jsonify({"task_uid": job_id})
@@ -167,10 +170,12 @@ def process_logs(path: ProcessLogsRequestPath):
     log_type = path.log_type
     seqid = path.seqid
     analysis = get_analysis_data(task_uid)
-    index_path = analysis.get_log_index(f"{log_type}.{seqid}.json")
+    index_path = analysis.get_log_index()
     if not index_path.exists():
         return dict(error="Data not found"), 404
-    index = orjson.loads(index_path.read_text())
+    index = get_log_index_for_process(index_path, seqid, log_type)
+    if not index:
+        return dict(error="Data not found"), 404
     blocks = index["blocks"]
     filter_values = request.args.getlist("filter[]")
     if filter_values:

requirements.txt

@@ -19,4 +19,6 @@ Pillow==10.4.0
 Perception==0.6.8
 # Perception missing peer dependency
 opencv-contrib-python-headless==4.11.0.86
+# Also peer vivisect dependency, but we use it directly
+msgpack==1.0.8
 boto3==1.38.15