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 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()})

View File

@ -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()

View File

@ -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 == []