From 4254f8cd375b254b6ed2aca34debcb505a94291b Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 9 Jul 2021 11:33:52 +0300 Subject: [PATCH] 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 --- .../services/ransomware/ransomware_report.py | 68 +++++++------------ .../ransomware/test_ransomware_report.py | 20 +++--- 2 files changed, 33 insertions(+), 55 deletions(-) diff --git a/monkey/monkey_island/cc/services/ransomware/ransomware_report.py b/monkey/monkey_island/cc/services/ransomware/ransomware_report.py index 769fc5e53..305fab1b6 100644 --- a/monkey/monkey_island/cc/services/ransomware/ransomware_report.py +++ b/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() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/ransomware/test_ransomware_report.py b/monkey/tests/unit_tests/monkey_island/cc/services/ransomware/test_ransomware_report.py index c1213ef95..3dbc9732f 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/ransomware/test_ransomware_report.py +++ b/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)