From 07616e483d107489f5a8e1b030a82f061d2aaee0 Mon Sep 17 00:00:00 2001 From: Prateek <126278220+kharanpv@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:33:57 -0400 Subject: [PATCH] perf: index MongoDB calls by task ID (#3125) * perf: index MongoDB calls by task ID * Update mongodb.py --- dev_utils/mongodb.py | 29 ++++++++++++++++++++ lib/cuckoo/core/startup.py | 17 +++++++----- tests/test_process_mongodb_startup.py | 38 ++++++++++++++++++++++++++ tests/test_startup_mongodb_indexes.py | 39 +++++++++++++++++++++++++++ utils/process.py | 3 ++- 5 files changed, 119 insertions(+), 7 deletions(-) create mode 100644 tests/test_process_mongodb_startup.py create mode 100644 tests/test_startup_mongodb_indexes.py diff --git a/dev_utils/mongodb.py b/dev_utils/mongodb.py index 1ddf4cc1b28..8705057a559 100644 --- a/dev_utils/mongodb.py +++ b/dev_utils/mongodb.py @@ -84,6 +84,35 @@ def files(self): return get_mongodb().files def __getattr__(self, name): return getattr(get_mongodb(), name) results_db = LegacyDB() +else: + class AutoReconnect(Exception): + pass + + class OperationFailure(Exception): + pass + + def connect_to_mongo(): + return None + + def get_mongodb(): + class DummyMongo: + @property + def client(self): + class DummyClient: + @property + def admin(self): + class DummyAdmin: + def command(self, *args, **kwargs): + raise OperationFailure("MongoDB is disabled") + return DummyAdmin() + return DummyClient() + return DummyMongo() + + class LegacyDB: + def __getattr__(self, name): + raise AttributeError("MongoDB is disabled") + + results_db = LegacyDB() MAX_AUTO_RECONNECT_ATTEMPTS = 5 diff --git a/lib/cuckoo/core/startup.py b/lib/cuckoo/core/startup.py index 3b9edad0388..1c7c94e8a89 100644 --- a/lib/cuckoo/core/startup.py +++ b/lib/cuckoo/core/startup.py @@ -304,21 +304,26 @@ def check_working_directory(): raise CuckooStartupError(f"Fix permission on tmpfs path: chown cape:cape {cuckoo.tmpfs.path}") -def check_webgui_mongo(): +def check_webgui_mongo(exit_on_connection_failure=True): if repconf.mongodb.enabled: from dev_utils.mongodb import connect_to_mongo, mongo_create_index client = connect_to_mongo() if not client: - sys.exit( - "You have enabled webgui but mongo isn't working, see mongodb manual for correct installation and configuration\nrun `systemctl status mongodb` for more info" + message = ( + "You have enabled webgui but mongo isn't working, see mongodb manual for correct installation and configuration\n" + "run `systemctl status mongodb` for more info" ) + if exit_on_connection_failure: + sys.exit(message) + log.warning(message) + return - # Create an index based on the info.id dict key. Increases overall scalability - # with large amounts of data. - # Note: Silently ignores the creation if the index already exists. + # Create required indexes to improve scalability with large amounts of data. + # Note: Silently ignores creation if an equivalent index already exists. index_configs = [ ("analysis", [("info.id", -1)], {"name": "info_id_desc"}), + ("calls", [("task_id", 1)], {"name": "task_id_1"}), ("files", [("_task_ids", 1)], {}), ] if repconf.mongodb.get("index_yara", False): diff --git a/tests/test_process_mongodb_startup.py b/tests/test_process_mongodb_startup.py new file mode 100644 index 00000000000..fc92280793e --- /dev/null +++ b/tests/test_process_mongodb_startup.py @@ -0,0 +1,38 @@ +import sys +from unittest.mock import Mock + +import pytest + +import utils.process as process + + +def test_main_reconciles_mongodb_indexes_before_modules(monkeypatch): + calls = [] + + monkeypatch.setattr(sys, "argv", ["process.py", "auto"]) + monkeypatch.setattr(process, "init_database", Mock()) + monkeypatch.setattr(process, "init_logging", Mock(return_value=[])) + + monkeypatch.setattr( + process, + "check_webgui_mongo", + Mock(side_effect=lambda **kwargs: calls.append(("mongo", kwargs))), + ) + monkeypatch.setattr( + process, + "init_modules", + Mock(side_effect=lambda: calls.append(("modules", {}))), + ) + monkeypatch.setattr( + process, + "autoprocess", + Mock(side_effect=RuntimeError("stop after startup")), + ) + + with pytest.raises(RuntimeError, match="stop after startup"): + process.main() + + assert calls == [ + ("mongo", {"exit_on_connection_failure": False}), + ("modules", {}), + ] diff --git a/tests/test_startup_mongodb_indexes.py b/tests/test_startup_mongodb_indexes.py new file mode 100644 index 00000000000..a0c6ecb97fb --- /dev/null +++ b/tests/test_startup_mongodb_indexes.py @@ -0,0 +1,39 @@ +from unittest.mock import Mock + +import lib.cuckoo.core.startup as startup + + +def test_check_webgui_mongo_creates_calls_task_id_index(monkeypatch): + create_index = Mock() + + monkeypatch.setattr(startup.repconf.mongodb, "enabled", True) + monkeypatch.setattr(startup.repconf.elasticsearchdb, "enabled", False) + + monkeypatch.setattr( + "dev_utils.mongodb.connect_to_mongo", + Mock(return_value=object()), + ) + monkeypatch.setattr( + "dev_utils.mongodb.mongo_create_index", + create_index, + ) + + startup.check_webgui_mongo() + + create_index.assert_any_call( + "calls", + [("task_id", 1)], + name="task_id_1", + ) + + +def test_check_webgui_mongo_can_fail_nonfatally(monkeypatch, caplog): + monkeypatch.setattr(startup.repconf.mongodb, "enabled", True) + monkeypatch.setattr( + "dev_utils.mongodb.connect_to_mongo", + Mock(return_value=None), + ) + + startup.check_webgui_mongo(exit_on_connection_failure=False) + + assert "mongo isn't working" in caplog.text diff --git a/utils/process.py b/utils/process.py index 64f1853268d..111dc9a6df0 100644 --- a/utils/process.py +++ b/utils/process.py @@ -45,7 +45,7 @@ TASK_REPORTED ) from lib.cuckoo.core.plugins import RunProcessing, RunReporting, RunSignatures -from lib.cuckoo.core.startup import ConsoleHandler, check_linux_dist, init_modules +from lib.cuckoo.core.startup import ConsoleHandler, check_linux_dist, check_webgui_mongo, init_modules cfg = Config() logconf = Config("logging") @@ -654,6 +654,7 @@ def main(): init_database() handlers = init_logging(debug=args.debug) + check_webgui_mongo(exit_on_connection_failure=False) init_modules() if args.id == "auto": autoprocess(