From 2bcf3b0a90b8ab5fecc1f032a9390238ec070f26 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 9 Jul 2021 16:03:16 +0300 Subject: [PATCH] Refactor ransomware report unit tests to mock "get_exploited()" method used. Also, minor refactorings in ransomware_report service and resource --- .../cc/resources/ransomware_report.py | 6 +- .../services/ransomware/ransomware_report.py | 110 +++++++++--------- .../ransomware/test_ransomware_report.py | 50 +++++--- 3 files changed, 89 insertions(+), 77 deletions(-) diff --git a/monkey/monkey_island/cc/resources/ransomware_report.py b/monkey/monkey_island/cc/resources/ransomware_report.py index a1b4c8ee8..0fb33e897 100644 --- a/monkey/monkey_island/cc/resources/ransomware_report.py +++ b/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()}) diff --git a/monkey/monkey_island/cc/services/ransomware/ransomware_report.py b/monkey/monkey_island/cc/services/ransomware/ransomware_report.py index 305fab1b6..ff8555cd3 100644 --- a/monkey/monkey_island/cc/services/ransomware/ransomware_report.py +++ b/monkey/monkey_island/cc/services/ransomware/ransomware_report.py @@ -5,64 +5,62 @@ 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"}}, - {"$unwind": "$data.files"}, - { - "$group": { - "_id": {"monkey_guid": "$monkey_guid", "files_encrypted": "$data.files.success"} - } - }, - {"$replaceRoot": {"newRoot": "$_id"}}, - {"$sort": {"files_encrypted": -1}}, - { - "$group": { - "_id": {"monkey_guid": "$monkey_guid"}, - "monkey_guid": {"$first": "$monkey_guid"}, - "files_encrypted": {"$first": "$files_encrypted"}, - } - }, - { - "$lookup": { - "from": "monkey", - "localField": "_id.monkey_guid", - "foreignField": "guid", - "as": "monkey", - } - }, - { - "$project": { - "monkey": {"$arrayElemAt": ["$monkey", 0]}, - "files_encrypted": "$files_encrypted", - } - }, - ] +def get_encrypted_files_table(): + query = [ + {"$match": {"telem_category": "file_encryption"}}, + {"$unwind": "$data.files"}, + { + "$group": { + "_id": {"monkey_guid": "$monkey_guid", "files_encrypted": "$data.files.success"} + } + }, + {"$replaceRoot": {"newRoot": "$_id"}}, + {"$sort": {"files_encrypted": -1}}, + { + "$group": { + "_id": {"monkey_guid": "$monkey_guid"}, + "monkey_guid": {"$first": "$monkey_guid"}, + "files_encrypted": {"$first": "$files_encrypted"}, + } + }, + { + "$lookup": { + "from": "monkey", + "localField": "_id.monkey_guid", + "foreignField": "guid", + "as": "monkey", + } + }, + { + "$project": { + "monkey": {"$arrayElemAt": ["$monkey", 0]}, + "files_encrypted": "$files_encrypted", + } + }, + ] - 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 + monkeys = list(mongo.db.telemetry.aggregate(query)) + exploited_nodes = ReportService.get_exploited() + for monkey in monkeys: + monkey["exploits"] = _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_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 3dbc9732f..ea9c4f293 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,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 = get_connection() + mongo = mongomock.MongoClient() 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": "WinDev2010Eval", "exploit": "SMB", "files_encrypted": True}, + {"hostname": "test-pc-2", "exploits": ["Manual execution"], "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 == []