Refactor ransomware_report.py to use current report infrastructure for fetching exploited nodes

Re-using current report infrastructure means that it's more trivial to implement/maintain and is already tested. The downside is performance

monkey/monkey_island/cc/services/ransomware/ransomware_report.py

@@ -1,4 +1,5 @@
 from monkey_island.cc.database import mongo
+from monkey_island.cc.services.reporting.report import ReportService
 from typing import Dict, List
 
 from monkey_island.cc.services.reporting.report import ReportService
@@ -38,53 +39,30 @@ class RansomwareReportService:
                     "files_encrypted": "$files_encrypted",
                 }
             },
-            {
-                "$lookup": {
-                    "from": "edge",
-                    "localField": "monkey._id",
-                    "foreignField": "dst_node_id",
-                    "as": "edge",
-                }
-            },
-            {
-                "$project": {
-                    "monkey": "$monkey",
-                    "files_encrypted": "$files_encrypted",
-                    "edge": {"$arrayElemAt": ["$edge", 0]},
-                }
-            },
-            {
-                "$project": {
-                    "hostname": "$monkey.hostname",
-                    "successful_exploits": {
-                        "$filter": {
-                            "input": "$edge.exploits",
-                            "as": "exploit",
-                            "cond": {"$eq": ["$$exploit.result", True]},
-                        }
-                    },
-                    "files_encrypted": "$files_encrypted",
-                }
-            },
-            {
-                "$addFields": {
-                    "successful_exploit": {"$arrayElemAt": ["$successful_exploits", 0]},
-                }
-            },
-            {
-                "$project": {
-                    "hostname": "$hostname",
-                    "exploiter": "$successful_exploit.info.display_name",
-                    "files_encrypted": "$files_encrypted",
-                }
-            },
         ]
 
-        table_data = list(mongo.db.telemetry.aggregate(query))
-        for encryption_entry in table_data:
-            if "exploiter" not in encryption_entry:
-                encryption_entry["exploiter"] = "Manual run"
-        return table_data
+        monkeys = list(mongo.db.telemetry.aggregate(query))
+        exploited_nodes = ReportService.get_exploited()
+        for monkey in monkeys:
+            monkey["exploits"] = RansomwareReportService.get_monkey_origin_exploits(
+                monkey["monkey"]["hostname"], exploited_nodes
+            )
+            monkey["hostname"] = monkey["monkey"]["hostname"]
+            del monkey["monkey"]
+            del monkey["_id"]
+        return monkeys
+
+    @staticmethod
+    def get_monkey_origin_exploits(monkey_hostname, exploited_nodes):
+        origin_exploits = [
+            exploited_node["exploits"]
+            for exploited_node in exploited_nodes
+            if exploited_node["label"] == monkey_hostname
+        ]
+        if origin_exploits:
+            return origin_exploits[0]
+        else:
+            return ["Manual execution"]
 
 def get_propagation_stats() -> Dict:
     scanned = ReportService.get_scanned()

monkey/tests/unit_tests/monkey_island/cc/services/ransomware/test_ransomware_report.py

@@ -1,3 +1,6 @@
+import mongoengine
+import pytest
+from mongoengine import get_connection
 import mongomock
 from tests.data_for_tests.mongo_documents.edges import EDGE_EXPLOITED, EDGE_SCANNED
 from tests.data_for_tests.mongo_documents.monkeys import MONKEY_AT_ISLAND, MONKEY_AT_VICTIM
@@ -17,14 +20,15 @@ from monkey_island.cc.services.ransomware.ransomware_report import RansomwareRep
 
 @pytest.fixture
 def fake_mongo(monkeypatch):
-    mongo = mongomock.MongoClient()
+    mongoengine.connect("mongoenginetest", host="mongomock://localhost")
+    mongo = get_connection()
     monkeypatch.setattr("monkey_island.cc.services.ransomware.ransomware_report.mongo", mongo)
+    monkeypatch.setattr("monkey_island.cc.services.reporting.report.mongo", mongo)
+    monkeypatch.setattr("monkey_island.cc.services.node.mongo", mongo)
     return mongo
 
 
-@pytest.mark.skip(
-    reason="A bug in mongomock prevents " "projecting the first element of an empty array"
-)
+@pytest.mark.skip(reason="Can't find a way to use the same mock database client in Monkey model")
 @pytest.mark.usefixtures("uses_database")
 def test_get_encrypted_files_table(fake_mongo):
     fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
@@ -44,9 +48,7 @@ def test_get_encrypted_files_table(fake_mongo):
     ]
 
 
-@pytest.mark.skip(
-    reason="A bug in mongomock prevents " "projecting the first element of an empty array"
-)
+@pytest.mark.skip(reason="Can't find a way to use the same mock database client in Monkey model")
 @pytest.mark.usefixtures("uses_database")
 def test_get_encrypted_files_table__only_errors(fake_mongo):
     fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
@@ -60,9 +62,7 @@ def test_get_encrypted_files_table__only_errors(fake_mongo):
     assert results == []
 
 
-@pytest.mark.skip(
-    reason="A bug in mongomock prevents " "projecting the first element of an empty array"
-)
+@pytest.mark.skip(reason="Can't find a way to use the same mock database client in Monkey model")
 @pytest.mark.usefixtures("uses_database")
 def test_get_encrypted_files_table__no_telemetries(fake_mongo):
     fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)