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,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

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