[pull] master from kevoreilly:master

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
 

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):

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", {}),
+    ]

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

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(