Alter usages of telemetry collection in report to store/fetch system info telemetry using the Telemetry model

This is required to automatically encrypt/decrypt the telemetries and it's a good practice to have a DAL for telemetries

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))}

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)
+        TelemetryModel.save_telemetry(telemetry_json)
-        return mongo.db.telemetry.find_one_or_404({"_id": telem_id})
+
+        return {}, 201
 
     @staticmethod
     def telemetry_to_displayed_telemetry(telemetry):

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},
         ):

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):
+def test_get_stolen_creds_no_creds(fake_mongo, uses_database):
-    fake_mongo.db.telemetry.insert_one(NO_CREDS_TELEMETRY_TELEM)
+    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 = []