Refactor ransomware_report.py to use current report infrastructure for fetching exploited nodes

Re-using current report infrastructure means that it's more trivial to implement/maintain and is already tested. The downside is performance
This commit is contained in:
VakarisZ 2021-07-09 11:33:52 +03:00
parent 9492b14c95
commit 4254f8cd37
2 changed files with 33 additions and 55 deletions

View File

@ -1,4 +1,5 @@
from monkey_island.cc.database import mongo
from monkey_island.cc.services.reporting.report import ReportService
from typing import Dict, List
from monkey_island.cc.services.reporting.report import ReportService
@ -38,53 +39,30 @@ class RansomwareReportService:
"files_encrypted": "$files_encrypted",
}
},
{
"$lookup": {
"from": "edge",
"localField": "monkey._id",
"foreignField": "dst_node_id",
"as": "edge",
}
},
{
"$project": {
"monkey": "$monkey",
"files_encrypted": "$files_encrypted",
"edge": {"$arrayElemAt": ["$edge", 0]},
}
},
{
"$project": {
"hostname": "$monkey.hostname",
"successful_exploits": {
"$filter": {
"input": "$edge.exploits",
"as": "exploit",
"cond": {"$eq": ["$$exploit.result", True]},
}
},
"files_encrypted": "$files_encrypted",
}
},
{
"$addFields": {
"successful_exploit": {"$arrayElemAt": ["$successful_exploits", 0]},
}
},
{
"$project": {
"hostname": "$hostname",
"exploiter": "$successful_exploit.info.display_name",
"files_encrypted": "$files_encrypted",
}
},
]
table_data = list(mongo.db.telemetry.aggregate(query))
for encryption_entry in table_data:
if "exploiter" not in encryption_entry:
encryption_entry["exploiter"] = "Manual run"
return table_data
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
@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_propagation_stats() -> Dict:
scanned = ReportService.get_scanned()

View File

@ -1,3 +1,6 @@
import mongoengine
import pytest
from mongoengine import get_connection
import mongomock
from tests.data_for_tests.mongo_documents.edges import EDGE_EXPLOITED, EDGE_SCANNED
from tests.data_for_tests.mongo_documents.monkeys import MONKEY_AT_ISLAND, MONKEY_AT_VICTIM
@ -17,14 +20,15 @@ from monkey_island.cc.services.ransomware.ransomware_report import RansomwareRep
@pytest.fixture
def fake_mongo(monkeypatch):
mongo = mongomock.MongoClient()
mongoengine.connect("mongoenginetest", host="mongomock://localhost")
mongo = get_connection()
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="A bug in mongomock prevents " "projecting the first element of an empty array"
)
@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):
fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
@ -44,9 +48,7 @@ def test_get_encrypted_files_table(fake_mongo):
]
@pytest.mark.skip(
reason="A bug in mongomock prevents " "projecting the first element of an empty array"
)
@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):
fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)
@ -60,9 +62,7 @@ def test_get_encrypted_files_table__only_errors(fake_mongo):
assert results == []
@pytest.mark.skip(
reason="A bug in mongomock prevents " "projecting the first element of an empty array"
)
@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):
fake_mongo.db.monkey.insert(MONKEY_AT_ISLAND)