Refactor ransomware report unit tests to mock "get_exploited()" method used. Also, minor refactorings in ransomware_report service and resource

monkey/monkey_island/cc/resources/ransomware_report.py

@@ -2,14 +2,12 @@ import flask_restful
 from flask import jsonify
 
 from monkey_island.cc.resources.auth.auth import jwt_required
-from monkey_island.cc.services.ransomware.ransomware_report import RansomwareReportService
 from monkey_island.cc.services.ransomware import ransomware_report
 
 
 class RansomwareReport(flask_restful.Resource):
     @jwt_required
     def get(self):
-        encrypted_files_table = RansomwareReportService.get_encrypted_files_table()
+        encrypted_files_table = ransomware_report.get_encrypted_files_table()
         return jsonify({"encrypted_files_table": encrypted_files_table,
-                        "propagation_stats": ransomware_report.get_propagation_stats()}
+                        "propagation_stats": ransomware_report.get_propagation_stats()})
-        )

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

@@ -5,8 +5,6 @@ from typing import Dict, List
 from monkey_island.cc.services.reporting.report import ReportService
 
 
-class RansomwareReportService:
-    @staticmethod
 def get_encrypted_files_table():
     query = [
         {"$match": {"telem_category": "file_encryption"}},
@@ -44,7 +42,7 @@ class RansomwareReportService:
     monkeys = list(mongo.db.telemetry.aggregate(query))
     exploited_nodes = ReportService.get_exploited()
     for monkey in monkeys:
-            monkey["exploits"] = RansomwareReportService.get_monkey_origin_exploits(
+        monkey["exploits"] = _get_monkey_origin_exploits(
             monkey["monkey"]["hostname"], exploited_nodes
         )
         monkey["hostname"] = monkey["monkey"]["hostname"]
@@ -52,8 +50,8 @@ class RansomwareReportService:
         del monkey["_id"]
     return monkeys
 
-    @staticmethod
+
-    def get_monkey_origin_exploits(monkey_hostname, exploited_nodes):
+def _get_monkey_origin_exploits(monkey_hostname, exploited_nodes):
     origin_exploits = [
         exploited_node["exploits"]
         for exploited_node in exploited_nodes

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

@@ -1,4 +1,4 @@
-import mongoengine
+import mongomock
 import pytest
 from mongoengine import get_connection
 import mongomock
@@ -15,22 +15,19 @@ import pytest
 from monkey_island.cc.services.ransomware import ransomware_report
 from monkey_island.cc.services.reporting.report import ReportService
 
-from monkey_island.cc.services.ransomware.ransomware_report import RansomwareReportService
+from monkey_island.cc.services.ransomware.ransomware_report import get_encrypted_files_table
+from monkey_island.cc.services.reporting.report import ReportService
 
 
 @pytest.fixture
 def fake_mongo(monkeypatch):
-    mongoengine.connect("mongoenginetest", host="mongomock://localhost")
+    mongo = mongomock.MongoClient()
-    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="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):
+def test_get_encrypted_files_table(fake_mongo, monkeypatch):
     fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
     fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM)
     fake_mongo.db.edge.insert(EDGE_EXPLOITED)
@@ -40,37 +37,56 @@ def test_get_encrypted_files_table(fake_mongo):
     fake_mongo.db.telemetry.insert(ENCRYPTION_ERROR)
     fake_mongo.db.telemetry.insert(ENCRYPTION_ONE_FILE)
 
-    results = RansomwareReportService.get_encrypted_files_table()
+    monkeypatch.setattr(
+        ReportService,
+        "get_exploited",
+        lambda: [{"label": "WinDev2010Eval", "exploits": ["SMB Exploiter"]}],
+    )
+
+    results = get_encrypted_files_table()
 
     assert results == [
-        {"hostname": "test-pc-2", "exploit": "Manual execution", "files_encrypted": True},
+        {"hostname": "test-pc-2", "exploits": ["Manual execution"], "files_encrypted": True},
-        {"hostname": "WinDev2010Eval", "exploit": "SMB", "files_encrypted": True},
+        {"hostname": "WinDev2010Eval", "exploits": ["SMB Exploiter"], "files_encrypted": True},
     ]
 
 
-@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):
+def test_get_encrypted_files_table__only_errors(fake_mongo, monkeypatch):
     fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
     fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM)
     fake_mongo.db.edge.insert(EDGE_EXPLOITED)
     fake_mongo.db.edge.insert(EDGE_SCANNED)
     fake_mongo.db.telemetry.insert(ENCRYPTION_ERROR)
 
-    results = RansomwareReportService.get_encrypted_files_table()
+    monkeypatch.setattr(
+        ReportService,
+        "get_exploited",
+        lambda: [{"label": "WinDev2010Eval", "exploits": ["SMB Exploiter"]}],
+    )
 
-    assert results == []
+    results = get_encrypted_files_table()
+
+    assert results == [
+        {"hostname": "test-pc-2", "exploits": ["Manual execution"], "files_encrypted": False}
+    ]
 
 
 @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):
+def test_get_encrypted_files_table__no_telemetries(fake_mongo, monkeypatch):
     fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
     fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM)
     fake_mongo.db.edge.insert(EDGE_EXPLOITED)
     fake_mongo.db.edge.insert(EDGE_SCANNED)
 
-    results = RansomwareReportService.get_encrypted_files_table()
+    monkeypatch.setattr(
+        ReportService,
+        "get_exploited",
+        lambda: [{"label": "WinDev2010Eval", "exploits": ["SMB Exploiter"]}],
+    )
+
+    results = get_encrypted_files_table()
 
     assert results == []