diff --git a/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py b/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py index 5573e5152..f784bc323 100644 --- a/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py +++ b/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py @@ -2,7 +2,7 @@ import flask_restful from bson import json_util from flask import request -from monkey_island.cc.database import mongo +from monkey_island.cc.models.telemetries import Telemetry from monkey_island.cc.resources.auth.auth import jwt_required @@ -10,4 +10,4 @@ class TelemetryBlackboxEndpoint(flask_restful.Resource): @jwt_required def get(self, **kw): find_query = json_util.loads(request.args.get("find_query")) - return {"results": list(mongo.db.telemetry.find(find_query))} + return {"results": list(Telemetry.get_telemetry_by_query(find_query))} diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py index 525197f0f..3ab2c1242 100644 --- a/monkey/monkey_island/cc/resources/telemetry.py +++ b/monkey/monkey_island/cc/resources/telemetry.py @@ -9,6 +9,7 @@ from flask import request from common.common_consts.telem_categories import TelemCategoryEnum from monkey_island.cc.database import mongo from monkey_island.cc.models.monkey import Monkey +from monkey_island.cc.models.telemetries.telemetry import Telemetry as TelemetryModel from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore from monkey_island.cc.services.node import NodeService @@ -37,7 +38,7 @@ class Telemetry(flask_restful.Resource): find_filter["timestamp"] = {"$gt": dateutil.parser.parse(timestamp)} result["objects"] = self.telemetry_to_displayed_telemetry( - mongo.db.telemetry.find(find_filter) + TelemetryModel.get_telemetry_by_query(query=find_filter) ) return result @@ -60,8 +61,9 @@ class Telemetry(flask_restful.Resource): process_telemetry(telemetry_json) - telem_id = mongo.db.telemetry.insert(telemetry_json) - return mongo.db.telemetry.find_one_or_404({"_id": telem_id}) + TelemetryModel.save_telemetry(telemetry_json) + + return {}, 201 @staticmethod def telemetry_to_displayed_telemetry(telemetry): diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py index 7d14d4f4a..32675f1d5 100644 --- a/monkey/monkey_island/cc/services/reporting/report.py +++ b/monkey/monkey_island/cc/services/reporting/report.py @@ -15,6 +15,7 @@ from common.network.network_range import NetworkRange from common.network.segmentation_utils import get_ip_in_src_and_not_in_dst from monkey_island.cc.database import mongo from monkey_island.cc.models import Monkey, Report +from monkey_island.cc.models.telemetries import Telemetry from monkey_island.cc.services.config import ConfigService from monkey_island.cc.services.configuration.utils import ( get_config_network_segments_as_subnet_groups, @@ -165,7 +166,7 @@ class ReportService: @staticmethod def _get_credentials_from_system_info_telems(): formatted_creds = [] - for telem in mongo.db.telemetry.find( + for telem in Telemetry.get_telemetry_by_query( {"telem_category": "system_info", "data.credentials": {"$exists": True}}, {"data.credentials": 1, "monkey_guid": 1}, ): diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py b/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py index 0093e4235..9f845b263 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py @@ -1,10 +1,11 @@ import datetime from copy import deepcopy -import mongomock +import mongoengine import pytest from bson import ObjectId +from monkey_island.cc.models.telemetries import Telemetry from monkey_island.cc.services.reporting.report import ReportService TELEM_ID = { @@ -49,6 +50,11 @@ SYSTEM_INFO_TELEMETRY_TELEM = { "_id": TELEM_ID["system_info_creds"], "monkey_guid": MONKEY_GUID, "telem_category": "system_info", + "timestamp": datetime.datetime(2021, 2, 19, 9, 0, 14, 984000), + "command_control_channel": { + "src": "192.168.56.1", + "dst": "192.168.56.2", + }, "data": { "credentials": { USER: { @@ -64,6 +70,11 @@ NO_CREDS_TELEMETRY_TELEM = { "_id": TELEM_ID["no_creds"], "monkey_guid": MONKEY_GUID, "telem_category": "exploit", + "timestamp": datetime.datetime(2021, 2, 19, 9, 0, 14, 984000), + "command_control_channel": { + "src": "192.168.56.1", + "dst": "192.168.56.2", + }, "data": { "machine": { "ip_addr": VICTIM_IP, @@ -125,13 +136,14 @@ NODE_DICT_FAILED_EXPLOITS["exploits"][1]["result"] = False @pytest.fixture def fake_mongo(monkeypatch): - mongo = mongomock.MongoClient() + mongo = mongoengine.connection.get_connection() monkeypatch.setattr("monkey_island.cc.services.reporting.report.mongo", mongo) + monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry.mongo", mongo) monkeypatch.setattr("monkey_island.cc.services.node.mongo", mongo) return mongo -def test_get_stolen_creds_exploit(fake_mongo): +def test_get_stolen_creds_exploit(fake_mongo, uses_database): fake_mongo.db.telemetry.insert_one(EXPLOIT_TELEMETRY_TELEM) stolen_creds_exploit = ReportService.get_stolen_creds() @@ -143,9 +155,9 @@ def test_get_stolen_creds_exploit(fake_mongo): assert expected_stolen_creds_exploit == stolen_creds_exploit -def test_get_stolen_creds_system_info(fake_mongo): +def test_get_stolen_creds_system_info(fake_mongo, uses_database): fake_mongo.db.monkey.insert_one(MONKEY_TELEM) - fake_mongo.db.telemetry.insert_one(SYSTEM_INFO_TELEMETRY_TELEM) + Telemetry.save_telemetry(SYSTEM_INFO_TELEMETRY_TELEM) stolen_creds_system_info = ReportService.get_stolen_creds() expected_stolen_creds_system_info = [ @@ -157,8 +169,9 @@ def test_get_stolen_creds_system_info(fake_mongo): assert expected_stolen_creds_system_info == stolen_creds_system_info -def test_get_stolen_creds_no_creds(fake_mongo): - fake_mongo.db.telemetry.insert_one(NO_CREDS_TELEMETRY_TELEM) +def test_get_stolen_creds_no_creds(fake_mongo, uses_database): + fake_mongo.db.monkey.insert_one(MONKEY_TELEM) + Telemetry.save_telemetry(NO_CREDS_TELEMETRY_TELEM) stolen_creds_no_creds = ReportService.get_stolen_creds() expected_stolen_creds_no_creds = []