From 549e621895ae0d92d4aba40e47b033109eb6dd88 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Fri, 4 Sep 2020 15:46:50 +0300
Subject: [PATCH] Small telemetry refactoring and added ScoutSuite telem

---
 monkey/infection_monkey/control.py            |  4 +--
 .../infection_monkey/telemetry/base_telem.py  | 36 +++++++++++++------
 .../telemetry/scoutsuite_telem.py             |  3 ++
 .../cc/resources/reporting/report.py          |  3 +-
 .../telemetry/processing/scoutsuite.py        | 11 ++++++
 5 files changed, 42 insertions(+), 15 deletions(-)
 create mode 100644 monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py

diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py
index 912514b8c..1f3c1c25e 100644
--- a/monkey/infection_monkey/control.py
+++ b/monkey/infection_monkey/control.py
@@ -128,12 +128,12 @@ class ControlClient(object):
             return {}
 
     @staticmethod
-    def send_telemetry(telem_category, data):
+    def send_telemetry(telem_category, json_data: str):
         if not WormConfiguration.current_server:
             LOG.error("Trying to send %s telemetry before current server is established, aborting." % telem_category)
             return
         try:
-            telemetry = {'monkey_guid': GUID, 'telem_category': telem_category, 'data': data}
+            telemetry = {'monkey_guid': GUID, 'telem_category': telem_category, 'data': json_data}
             requests.post("https://%s/api/telemetry" % (WormConfiguration.current_server,),  # noqa: DUO123
                           data=json.dumps(telemetry),
                           headers={'content-type': 'application/json'},
diff --git a/monkey/infection_monkey/telemetry/base_telem.py b/monkey/infection_monkey/telemetry/base_telem.py
index 7617ab4e3..07559491c 100644
--- a/monkey/infection_monkey/telemetry/base_telem.py
+++ b/monkey/infection_monkey/telemetry/base_telem.py
@@ -5,6 +5,7 @@ import logging
 from infection_monkey.control import ControlClient
 
 logger = logging.getLogger(__name__)
+LOGGED_DATA_LENGTH = 300 # How many characters of telemetry data will be logged
 
 __author__ = 'itay.mizeretz'
 
@@ -22,12 +23,25 @@ class BaseTelem(object, metaclass=abc.ABCMeta):
         Sends telemetry to island
         """
         data = self.get_data()
+        serialized_data = json.dumps(data, cls=self.json_encoder)
+        self.log_telem_sending(serialized_data, log_data)
+        ControlClient.send_telemetry(self.telem_category, serialized_data)
+
+    @abc.abstractmethod
+    def get_data(self) -> dict:
+        """
+        :return: Data of telemetry (should be dict)
+        """
+        pass
+
+    @property
+    def json_encoder(self):
+        return json.JSONEncoder
+
+    def log_telem_sending(self, serialized_data: str, log_data=True):
+        logger.debug(f"Sending {self.telem_category} telemetry.")
         if log_data:
-            data_to_log = json.dumps(data)
-        else:
-            data_to_log = 'redacted'
-        logger.debug("Sending {} telemetry. Data: {}".format(self.telem_category, data_to_log))
-        ControlClient.send_telemetry(self.telem_category, data)
+            logger.debug(f"Telemetry contents: {BaseTelem.truncate_data(serialized_data)}")
 
     @property
     @abc.abstractmethod
@@ -37,9 +51,9 @@ class BaseTelem(object, metaclass=abc.ABCMeta):
         """
         pass
 
-    @abc.abstractmethod
-    def get_data(self) -> dict:
-        """
-        :return: Data of telemetry (should be dict)
-        """
-        pass
+    @staticmethod
+    def truncate_data(data: str):
+        if len(data) <= LOGGED_DATA_LENGTH:
+            return data
+        else:
+            return f"{data[:LOGGED_DATA_LENGTH]}..."
diff --git a/monkey/infection_monkey/telemetry/scoutsuite_telem.py b/monkey/infection_monkey/telemetry/scoutsuite_telem.py
index d606ea3c3..743de93e7 100644
--- a/monkey/infection_monkey/telemetry/scoutsuite_telem.py
+++ b/monkey/infection_monkey/telemetry/scoutsuite_telem.py
@@ -1,4 +1,6 @@
 from common.common_consts.telem_categories import TelemCategoryEnum
+from infection_monkey.system_info.collectors.scoutsuite_collector.scoutsuite.ScoutSuite.output.result_encoder import \
+    ScoutJsonEncoder
 from infection_monkey.telemetry.base_telem import BaseTelem
 
 
@@ -12,6 +14,7 @@ class ScoutSuiteTelem(BaseTelem):
         super().__init__()
         self.data = data
 
+    json_encoder = ScoutJsonEncoder
     telem_category = TelemCategoryEnum.SCOUTSUITE
 
     def get_data(self):
diff --git a/monkey/monkey_island/cc/resources/reporting/report.py b/monkey/monkey_island/cc/resources/reporting/report.py
index a0ea8b0b9..f196fdfb6 100644
--- a/monkey/monkey_island/cc/resources/reporting/report.py
+++ b/monkey/monkey_island/cc/resources/reporting/report.py
@@ -31,8 +31,7 @@ class Report(flask_restful.Resource):
                     "statusesToPillars": ZeroTrustService.get_statuses_to_pillars(),
                     "pillarsToStatuses": ZeroTrustService.get_pillars_to_statuses(),
                     "grades": ZeroTrustService.get_pillars_grades()
-                }
-                )
+                })
             elif report_data == REPORT_DATA_PRINCIPLES_STATUS:
                 return jsonify(ZeroTrustService.get_principles_status())
             elif report_data == REPORT_DATA_FINDINGS:
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py b/monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py
new file mode 100644
index 000000000..ae63fe508
--- /dev/null
+++ b/monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py
@@ -0,0 +1,11 @@
+from monkey_island.cc.database import mongo
+
+
+def process_scoutsuite_telemetry(telemetry_json):
+    update_data(telemetry_json)
+
+
+def update_data(telemetry_json):
+    mongo.db.scoutsuite.update(
+        {'guid': telemetry_json['monkey_guid']},
+        {'$push': {'results': telemetry_json['data']}})