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