From 7538f774eda75ce3e129bc9400c8723aaab0cded Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Thu, 3 Sep 2020 12:38:40 +0300
Subject: [PATCH] Migrated more hard coded telem category values to use enum

---
 monkey/common/common_consts/telem_categories.py |  1 +
 .../scoutsuite_collector.py                     |  2 +-
 .../telemetry/attack/attack_telem.py            |  3 ++-
 .../services/telemetry/processing/processing.py | 17 +++++++++--------
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/monkey/common/common_consts/telem_categories.py b/monkey/common/common_consts/telem_categories.py
index c983786b9..70066d290 100644
--- a/monkey/common/common_consts/telem_categories.py
+++ b/monkey/common/common_consts/telem_categories.py
@@ -7,3 +7,4 @@ class TelemCategoryEnum:
     SYSTEM_INFO = 'system_info'
     TRACE = 'trace'
     TUNNEL = 'tunnel'
+    ATTACK = 'attack'
diff --git a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py b/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py
index fb33cce4b..f7d6b7ec5 100644
--- a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py
+++ b/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py
@@ -20,4 +20,4 @@ def run_scoutsuite(cloud_type):
 
 
 def send_results(results):
-    ScoutSuiteTelem.send(results)
+    ScoutSuiteTelem(results).send(results)
diff --git a/monkey/infection_monkey/telemetry/attack/attack_telem.py b/monkey/infection_monkey/telemetry/attack/attack_telem.py
index 893f4492a..ba3fae8fd 100644
--- a/monkey/infection_monkey/telemetry/attack/attack_telem.py
+++ b/monkey/infection_monkey/telemetry/attack/attack_telem.py
@@ -1,3 +1,4 @@
+from common.common_consts.telem_categories import TelemCategoryEnum
 from infection_monkey.telemetry.base_telem import BaseTelem
 
 __author__ = "VakarisZ"
@@ -15,7 +16,7 @@ class AttackTelem(BaseTelem):
         self.technique = technique
         self.status = status
 
-    telem_category = 'attack'
+    telem_category = TelemCategoryEnum.ATTACK
 
     def get_data(self):
         return {
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/processing.py b/monkey/monkey_island/cc/services/telemetry/processing/processing.py
index 566c11dcc..960a01517 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/processing.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/processing.py
@@ -1,5 +1,6 @@
 import logging
 
+from common.common_consts.telem_categories import TelemCategoryEnum
 from monkey_island.cc.services.telemetry.processing.exploit import \
     process_exploit_telemetry
 from monkey_island.cc.services.telemetry.processing.post_breach import \
@@ -17,15 +18,15 @@ logger = logging.getLogger(__name__)
 
 TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = \
     {
-        'tunnel': process_tunnel_telemetry,
-        'state': process_state_telemetry,
-        'exploit': process_exploit_telemetry,
-        'scan': process_scan_telemetry,
-        'system_info': process_system_info_telemetry,
-        'post_breach': process_post_breach_telemetry,
+        TelemCategoryEnum.TUNNEL: process_tunnel_telemetry,
+        TelemCategoryEnum.STATE: process_state_telemetry,
+        TelemCategoryEnum.EXPLOIT: process_exploit_telemetry,
+        TelemCategoryEnum.SCAN: process_scan_telemetry,
+        TelemCategoryEnum.SYSTEM_INFO: process_system_info_telemetry,
+        TelemCategoryEnum.POST_BREACH: process_post_breach_telemetry,
         # `lambda *args, **kwargs: None` is a no-op.
-        'trace': lambda *args, **kwargs: None,
-        'attack': lambda *args, **kwargs: None,
+        TelemCategoryEnum.TRACE: lambda *args, **kwargs: None,
+        TelemCategoryEnum.ATTACK: lambda *args, **kwargs: None,
     }