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

This commit is contained in:
VakarisZ 2021-07-09 16:03:16 +03:00
parent 4254f8cd37
commit 2bcf3b0a90
3 changed files with 89 additions and 77 deletions

View File

@ -2,14 +2,12 @@ import flask_restful
from flask import jsonify from flask import jsonify
from monkey_island.cc.resources.auth.auth import jwt_required 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 from monkey_island.cc.services.ransomware import ransomware_report
class RansomwareReport(flask_restful.Resource): class RansomwareReport(flask_restful.Resource):
@jwt_required @jwt_required
def get(self): 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, return jsonify({"encrypted_files_table": encrypted_files_table,
"propagation_stats": ransomware_report.get_propagation_stats()} "propagation_stats": ransomware_report.get_propagation_stats()})
)

View File

@ -5,64 +5,62 @@ from typing import Dict, List
from monkey_island.cc.services.reporting.report import ReportService from monkey_island.cc.services.reporting.report import ReportService
class RansomwareReportService: def get_encrypted_files_table():
@staticmethod query = [
def get_encrypted_files_table(): {"$match": {"telem_category": "file_encryption"}},
query = [ {"$unwind": "$data.files"},
{"$match": {"telem_category": "file_encryption"}}, {
{"$unwind": "$data.files"}, "$group": {
{ "_id": {"monkey_guid": "$monkey_guid", "files_encrypted": "$data.files.success"}
"$group": { }
"_id": {"monkey_guid": "$monkey_guid", "files_encrypted": "$data.files.success"} },
} {"$replaceRoot": {"newRoot": "$_id"}},
}, {"$sort": {"files_encrypted": -1}},
{"$replaceRoot": {"newRoot": "$_id"}}, {
{"$sort": {"files_encrypted": -1}}, "$group": {
{ "_id": {"monkey_guid": "$monkey_guid"},
"$group": { "monkey_guid": {"$first": "$monkey_guid"},
"_id": {"monkey_guid": "$monkey_guid"}, "files_encrypted": {"$first": "$files_encrypted"},
"monkey_guid": {"$first": "$monkey_guid"}, }
"files_encrypted": {"$first": "$files_encrypted"}, },
} {
}, "$lookup": {
{ "from": "monkey",
"$lookup": { "localField": "_id.monkey_guid",
"from": "monkey", "foreignField": "guid",
"localField": "_id.monkey_guid", "as": "monkey",
"foreignField": "guid", }
"as": "monkey", },
} {
}, "$project": {
{ "monkey": {"$arrayElemAt": ["$monkey", 0]},
"$project": { "files_encrypted": "$files_encrypted",
"monkey": {"$arrayElemAt": ["$monkey", 0]}, }
"files_encrypted": "$files_encrypted", },
} ]
},
]
monkeys = list(mongo.db.telemetry.aggregate(query)) monkeys = list(mongo.db.telemetry.aggregate(query))
exploited_nodes = ReportService.get_exploited() exploited_nodes = ReportService.get_exploited()
for monkey in monkeys: for monkey in monkeys:
monkey["exploits"] = RansomwareReportService.get_monkey_origin_exploits( monkey["exploits"] = _get_monkey_origin_exploits(
monkey["monkey"]["hostname"], exploited_nodes monkey["monkey"]["hostname"], exploited_nodes
) )
monkey["hostname"] = monkey["monkey"]["hostname"] monkey["hostname"] = monkey["monkey"]["hostname"]
del monkey["monkey"] del monkey["monkey"]
del monkey["_id"] del monkey["_id"]
return monkeys return monkeys
@staticmethod
def get_monkey_origin_exploits(monkey_hostname, exploited_nodes): def _get_monkey_origin_exploits(monkey_hostname, exploited_nodes):
origin_exploits = [ origin_exploits = [
exploited_node["exploits"] exploited_node["exploits"]
for exploited_node in exploited_nodes for exploited_node in exploited_nodes
if exploited_node["label"] == monkey_hostname if exploited_node["label"] == monkey_hostname
] ]
if origin_exploits: if origin_exploits:
return origin_exploits[0] return origin_exploits[0]
else: else:
return ["Manual execution"] return ["Manual execution"]
def get_propagation_stats() -> Dict: def get_propagation_stats() -> Dict:
scanned = ReportService.get_scanned() scanned = ReportService.get_scanned()

View File

@ -1,4 +1,4 @@
import mongoengine import mongomock
import pytest import pytest
from mongoengine import get_connection from mongoengine import get_connection
import mongomock import mongomock
@ -15,22 +15,19 @@ import pytest
from monkey_island.cc.services.ransomware import ransomware_report from monkey_island.cc.services.ransomware import ransomware_report
from monkey_island.cc.services.reporting.report import ReportService 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 @pytest.fixture
def fake_mongo(monkeypatch): 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.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 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") @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_ISLAND)
fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM) fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM)
fake_mongo.db.edge.insert(EDGE_EXPLOITED) 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_ERROR)
fake_mongo.db.telemetry.insert(ENCRYPTION_ONE_FILE) 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 == [ 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") @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_ISLAND)
fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM) fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM)
fake_mongo.db.edge.insert(EDGE_EXPLOITED) fake_mongo.db.edge.insert(EDGE_EXPLOITED)
fake_mongo.db.edge.insert(EDGE_SCANNED) fake_mongo.db.edge.insert(EDGE_SCANNED)
fake_mongo.db.telemetry.insert(ENCRYPTION_ERROR) 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.skip(reason="Can't find a way to use the same mock database client in Monkey model")
@pytest.mark.usefixtures("uses_database") @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_ISLAND)
fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM) fake_mongo.db.monkey.insert(MONKEY_AT_VICTIM)
fake_mongo.db.edge.insert(EDGE_EXPLOITED) fake_mongo.db.edge.insert(EDGE_EXPLOITED)
fake_mongo.db.edge.insert(EDGE_SCANNED) 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 == [] assert results == []