Improve performance and error handling in bson part 2 (kevoreilly#2814)

* Improve performance and error handling in core modules

Refactored compressor and behavior processing to use mmap for faster file access, added standard file reading fallback, and improved error handling and logging in guest and sniffer modules. Refactored behavior summary logic to use dispatch tables for API handling, enhancing maintainability and extensibility. Enhanced remote sniffer operations with timeouts and error logging for robustness.

* sync

* Apply suggestions from code review

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update modules/auxiliary/sniffer.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update compressor.py

* Update compressor.py

* Update compressor.py

* Update compressor.py

* fix

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: doomedraven

conf/default/processing.conf.default

@@ -54,6 +54,9 @@ ram_boost = no
 replace_patterns = no
 file_activities = no
 
+# process behavior files in ram to speedup processing a little bit?
+ram_mmap = no
+
 [tracee]
 enabled = no
 

lib/cuckoo/common/compressor.py

@@ -1,5 +1,6 @@
 import binascii
 import logging
+import mmap
 import os
 import struct
 from pathlib import Path
@@ -76,50 +77,115 @@ def __init__(self):
         self.callmap = {}
         self.head = []
         self.ccounter = 0
+        self.category = None
 
-    def __next_message(self):
-        data = self.fd_in.read(4)
-        if not data:
-            return (False, False)
-        _size = struct.unpack("I", data)[0]
-        data += self.fd_in.read(_size - 4)
-        self.raw_data = data
-        return (data, bson.decode(data))
-
-    def run(self, file_path):
-        if not os.path.isfile(file_path) and os.stat(file_path).st_size:
+    def _process_message(self, msg, data):
+        mtype = msg.get("type")  # message type [debug, new_process, info]
+        if mtype in {"debug", "new_process", "info"}:
+            self.category = msg.get("category", "None")
+            self.head.append(data.tobytes() if isinstance(data, memoryview) else data)
+
+        elif self.category and self.category.startswith("__"):
+            self.head.append(data.tobytes() if isinstance(data, memoryview) else data)
+        else:
+            tid = msg.get("T", -1)
+            time = msg.get("t", 0)
+
+            if tid not in self.threads:
+                self.threads[tid] = Compressor(100)
+
+            csum = self.checksum(msg)
+            self.ccounter += 1
+            v = (csum, self.ccounter, time)
+            self.threads[tid].add(v)
+
+            if csum not in self.callmap:
+                self.callmap[csum] = msg
+
+    def _process_mmap_content(self, mm):
+        with memoryview(mm) as mv:
+            offset = 0
+            size_mm = len(mm)
+
+            while offset < size_mm:
+                # Read size (4 bytes)
+                if offset + 4 > size_mm:
+                    break
+
+                # Slicing memoryview returns memoryview
+                size_bytes = mv[offset : offset + 4]
+                _size = struct.unpack("I", size_bytes)[0]
+
+                if offset + _size > size_mm:
+                    break
+
+                data = mv[offset : offset + _size]
+                offset += _size
+
+                try:
+                    msg = bson.decode(data)
+                except Exception:
+                    break
+
+                if msg:
+                    self._process_message(msg, data)
+
+    def run(self, file_path, use_mmap=False):
+        if use_mmap:
+            return self._run_mmap(file_path)
+        return self._run_standard(file_path)
+
+    def _run_mmap(self, file_path):
+        if not os.path.isfile(file_path) or not os.stat(file_path).st_size:
             log.warning("File %s does not exists or it is invalid", file_path)
             return False
 
-        self.fd_in = open(file_path, "rb")
+        with open(file_path, "rb") as f:
+            try:
+                mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
+            except ValueError:
+                return False
 
-        msg = "---"
-        while msg:
-            data, msg = self.__next_message()
+            try:
+                self._process_mmap_content(mm)
+            finally:
+                mm.close()
 
-            if msg:
-                mtype = msg.get("type")  # message type [debug, new_process, info]
-                if mtype in {"debug", "new_process", "info"}:
-                    self.category = msg.get("category", "None")
-                    self.head.append(data)
+        return self.flush(file_path)
 
-                elif self.category.startswith("__"):
-                    self.head.append(data)
-                else:
-                    tid = msg.get("T", -1)
-                    time = msg.get("t", 0)
+    def _run_standard(self, file_path):
+        if not os.path.isfile(file_path) or not os.stat(file_path).st_size:
+            log.warning("File %s does not exists or it is invalid", file_path)
+            return False
+
+        with open(file_path, "rb") as f:
+            while True:
+                size_bytes = f.read(4)
+                if len(size_bytes) < 4:
+                    break
+
+                try:
+                    _size = struct.unpack("I", size_bytes)[0]
+                except struct.error:
+                    break
 
-                    if tid not in self.threads:
-                        self.threads[tid] = Compressor(100)
+                remaining = _size - 4
+                if remaining < 0:
+                    break
 
-                    csum = self.checksum(msg)
-                    self.ccounter += 1
-                    v = (csum, self.ccounter, time)
-                    self.threads[tid].add(v)
+                data_body = f.read(remaining)
+                if len(data_body) < remaining:
+                    break
 
-                    if csum not in self.callmap:
-                        self.callmap[csum] = msg
-        self.fd_in.close()
+                data = size_bytes + data_body
+
+                try:
+                    msg = bson.decode(data)
+                except Exception:
+                    break
+
+                if msg:
+                    self._process_message(msg, data)
 
         return self.flush(file_path)
 
@@ -171,3 +237,31 @@ def checksum(self, msg):
         content = f"{index}{msg['T']}{msg['R']}{args}{self.category}{msg['P']}"
 
         return binascii.crc32(content.encode("utf8"))
+
+
+if __name__ == "__main__":
+    import argparse
+    import time
+    import sys
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file", help="Path to BSON file to compress")
+    parser.add_argument("--mmap", action="store_true", help="Use mmap for compression")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.file):
+        print(f"File {args.file} not found.")
+        sys.exit(1)
+
+    print(f"Compressing {args.file}...")
+    start = time.time()
+
+    compressor = CuckooBsonCompressor()
+    result = compressor.run(args.file, use_mmap=args.mmap)
+
+    end = time.time()
+
+    if result:
+        print(f"Compression successful. Took {end - start:.4f} seconds.")
+    else:
+        print("Compression failed.")

lib/cuckoo/core/guest.py

@@ -107,11 +107,11 @@ def get(self, method, *args, **kwargs):
 
             try:
                 r = session.get(url, *args, **kwargs)
-            except requests.ConnectionError:
+            except requests.ConnectionError as e:
                 raise CuckooGuestError(
-                    "CAPE Agent failed without error status, please try "
-                    "upgrading to the latest version of agent.py (>= 0.10) and "
-                    "notify us if the issue persists"
+                    f"CAPE Agent failed without error status, please try "
+                    f"upgrading to the latest version of agent.py (>= 0.10) and "
+                    f"notify us if the issue persists. Error: {e}"
                 )
 
         do_raise and r.raise_for_status()
@@ -134,11 +134,11 @@ def post(self, method, *args, **kwargs):
 
         try:
             r = session.post(url, *args, **kwargs)
-        except requests.ConnectionError:
+        except requests.ConnectionError as e:
             raise CuckooGuestError(
-                "CAPE Agent failed without error status, please try "
-                "upgrading to the latest version of agent.py (>= 0.10) and "
-                "notify us if the issue persists"
+                f"CAPE Agent failed without error status, please try "
+                f"upgrading to the latest version of agent.py (>= 0.10) and "
+                f"notify us if the issue persists. Error: {e}"
             )
 
         r.raise_for_status()
@@ -377,14 +377,15 @@ def wait_for_completion(self):
 
             try:
                 status = self.get("/status", timeout=5).json()
-            except (CuckooGuestError, requests.exceptions.ReadTimeout):
+            except (CuckooGuestError, requests.exceptions.ReadTimeout) as e:
                 # this might fail due to timeouts or just temporary network
                 # issues thus we don't want to abort the analysis just yet and
                 # wait for things to recover
                 log.warning(
-                    "Task #%s: Virtual Machine %s /status failed. This can indicate the guest losing network connectivity",
+                    "Task #%s: Virtual Machine %s /status failed. This can indicate the guest losing network connectivity. Error: %s",
                     self.task_id,
                     self.vmid,
+                    e,
                 )
                 continue
             except Exception as e:

modules/auxiliary/sniffer.py

@@ -191,27 +191,43 @@ def start(self):
                 f.write(f"echo $PID > /tmp/{self.task.id}.pid")
                 f.write("\n")
 
-            subprocess.check_output(
-                ["scp", "-q", f"/tmp/{self.task.id}.sh", remote_host + f":/tmp/{self.task.id}.sh"],
-            )
-            subprocess.check_output(
-                ["ssh", remote_host, "nohup", "/bin/bash", f"/tmp/{self.task.id}.sh", ">", "/tmp/log", "2>", "/tmp/err"],
-            )
-
-            self.pid = subprocess.check_output(
-                ["ssh", remote_host, "cat", f"/tmp/{self.task.id}.pid"], stderr=subprocess.DEVNULL
-            ).strip()
-            log.info(
-                "Started remote sniffer @ %s with (interface=%s, host=%s, dump path=%s, pid=%s)",
-                remote_host,
-                interface,
-                host,
-                file_path,
-                self.pid,
-            )
-            subprocess.check_output(
-                ["ssh", remote_host, "rm", "-f", f"/tmp/{self.task.id}.pid", f"/tmp/{self.task.id}.sh"],
-            )
+            try:
+                subprocess.check_output(
+                    ["scp", "-q", f"/tmp/{self.task.id}.sh", remote_host + f":/tmp/{self.task.id}.sh"], timeout=30
+                )
+                subprocess.check_output(
+                    [
+                        "ssh",
+                        remote_host,
+                        "nohup",
+                        "/bin/bash",
+                        f"/tmp/{self.task.id}.sh",
+                        ">",
+                        "/tmp/log",
+                        "2>",
+                        "/tmp/err",
+                    ],
+                    timeout=30,
+                )
+
+                self.pid = subprocess.check_output(
+                    ["ssh", remote_host, "cat", f"/tmp/{self.task.id}.pid"], stderr=subprocess.DEVNULL, timeout=30
+                ).strip()
+                log.info(
+                    "Started remote sniffer @ %s with (interface=%s, host=%s, dump path=%s, pid=%s)",
+                    remote_host,
+                    interface,
+                    host,
+                    file_path,
+                    self.pid,
+                )
+                subprocess.check_output(
+                    ["ssh", remote_host, "rm", "-f", f"/tmp/{self.task.id}.pid", f"/tmp/{self.task.id}.sh"], timeout=30
+                )
+            except subprocess.TimeoutExpired:
+                log.error("Timeout connecting to remote host %s", remote_host)
+            except subprocess.CalledProcessError as e:
+                log.error("Error connecting to remote host %s: %s", remote_host, e)
 
         else:
             try:
@@ -220,7 +236,13 @@ def start(self):
                 log.exception("Failed to start sniffer (interface=%s, host=%s, dump path=%s)", interface, host, file_path)
                 return
 
-            log.info("Started sniffer with PID %d (interface=%s, host=%s, dump path=%s)", self.proc.pid, interface, host, file_path)
+            log.info(
+                "Started sniffer with PID %d (interface=%s, host=%s, dump path=%s)",
+                self.proc.pid,
+                interface,
+                host,
+                file_path,
+            )
 
     def stop(self):
         """Stop sniffing.
@@ -235,22 +257,26 @@ def stop(self):
             remote_host = self.options.get("host", "")
             remote_args = ["ssh", remote_host, "kill", "-2", self.pid]
 
-            subprocess.check_output(remote_args)
+            try:
+                subprocess.check_output(remote_args, timeout=30)
 
-            file_path = os.path.join(CUCKOO_ROOT, "storage", "analyses", str(self.task.id), "dump.pcap")
-            file_path2 = f"/tmp/tcp.dump.{self.task.id}"
+                file_path = os.path.join(CUCKOO_ROOT, "storage", "analyses", str(self.task.id), "dump.pcap")
+                file_path2 = f"/tmp/tcp.dump.{self.task.id}"
 
-            subprocess.check_output(["scp", "-q", f"{remote_host}:{file_path2}", file_path])
-            subprocess.check_output(["ssh", remote_host, "rm", "-f", file_path2])
+                subprocess.check_output(["scp", "-q", f"{remote_host}:{file_path2}", file_path], timeout=300)
+                subprocess.check_output(["ssh", remote_host, "rm", "-f", file_path2], timeout=30)
+            except (subprocess.TimeoutExpired, subprocess.CalledProcessError) as e:
+                log.error("Error stopping remote sniffer: %s", e)
             return
 
         if self.proc and not self.proc.poll():
             if self.proc.args[0] == self.sudo_path and "-Z" in self.proc.args:
                 # We must kill the child process that sudo spawned. We won't
                 # have permission to kill the parent process because it's owned by root.
                 try:
-                    pid = int(subprocess.check_output(["ps", "--ppid", str(self.proc.pid), "-o", "pid="]).decode())
-                except (subprocess.CalledProcessError, TypeError, ValueError):
+                    output = subprocess.check_output(["ps", "--ppid", str(self.proc.pid), "-o", "pid="]).decode().strip()
+                    pid = int(output.split()[0])
+                except (subprocess.CalledProcessError, TypeError, ValueError, IndexError):
                     log.exception("Failed to get child pid of sudo process to stop the sniffer.")
                     return
                 term_func = functools.partial(os.kill, pid, signal.SIGTERM)
@@ -261,14 +287,14 @@ def stop(self):
                 pid = self.proc.pid
             try:
                 term_func()
-                _, _ = self.proc.communicate()
+                _, _ = self.proc.communicate(timeout=5)
             except Exception as e:
                 log.error("Unable to stop the sniffer (first try) with pid %d: %s", pid, e)
                 try:
                     if not self.proc.poll():
                         log.debug("Killing sniffer")
                         kill_func()
-                        _, _ = self.proc.communicate()
+                        _, _ = self.proc.communicate(timeout=5)
                 except OSError as e:
                     log.debug("Error killing sniffer: %s, continuing", e)
                 except Exception as e: