forked from p15670423/monkey
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:
parent
4254f8cd37
commit
2bcf3b0a90
|
@ -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()})
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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 == []
|
||||
|
||||
|
|
Loading…
Reference in New Issue