Exported UsageTechnique class to separate file, improved documentation. Refactored scripting attack telemetry sending in pba

monkey/infection_monkey/post_breach/pba.py

@@ -47,10 +47,21 @@ class PBA(object):
         """
         exec_funct = self._execute_default
         result = exec_funct()
-        if result[1] and isinstance(self.command, list) and len(self.command) > 1:
+        if self.scripts_were_used(result):
             T1064Telem(ScanStatus.USED, "Scripts used to execute %s post breach action." % self.name).send()
         PostBreachTelem(self, result).send()
 
+    def scripts_were_used(self, pba_execution_result):
+        """
+        Determines if scripts were used to execute PBA
+        :param pba_execution_result: result of execution function. e.g. self._execute_default
+        :return: True if scripts were used, False otherwise
+        """
+        pba_execution_succeeded = pba_execution_result[1]
+        if pba_execution_succeeded and isinstance(self.command, list) and len(self.command) > 1:
+            return True
+        return False
+
     def _execute_default(self):
         """
         Default post breach command execution routine

monkey/monkey_island/cc/services/attack/technique_reports/T1035.py

@@ -1,5 +1,4 @@
-from monkey_island.cc.database import mongo
-from monkey_island.cc.services.attack.technique_reports import UsageTechnique
+from monkey_island.cc.services.attack.technique_reports.usage_technique import UsageTechnique
 
 __author__ = "VakarisZ"
 

monkey/monkey_island/cc/services/attack/technique_reports/T1064.py

@@ -1,4 +1,4 @@
-from monkey_island.cc.services.attack.technique_reports import UsageTechnique
+from monkey_island.cc.services.attack.technique_reports.usage_technique import UsageTechnique
 
 __author__ = "VakarisZ"
 

monkey/monkey_island/cc/services/attack/technique_reports/T1106.py

@@ -1,4 +1,4 @@
-from monkey_island.cc.services.attack.technique_reports import UsageTechnique
+from monkey_island.cc.services.attack.technique_reports.usage_technique import UsageTechnique
 
 __author__ = "VakarisZ"
 

monkey/monkey_island/cc/services/attack/technique_reports/T1129.py

@@ -1,5 +1,4 @@
-from monkey_island.cc.database import mongo
-from monkey_island.cc.services.attack.technique_reports import UsageTechnique
+from monkey_island.cc.services.attack.technique_reports.usage_technique import UsageTechnique
 
 __author__ = "VakarisZ"
 

monkey/monkey_island/cc/services/attack/technique_reports/__init__.py

@@ -2,7 +2,7 @@ import abc
 import logging
 
 from monkey_island.cc.database import mongo
-from common.utils.attack_utils import ScanStatus, UsageEnum
+from common.utils.attack_utils import ScanStatus
 from monkey_island.cc.services.attack.attack_config import AttackConfig
 from common.utils.code_utils import abstractstatic
 
@@ -115,47 +115,3 @@ class AttackTechnique(object):
         data = cls.get_message_and_status(status)
         data.update({'title': cls.technique_title()})
         return data
-
-
-class UsageTechnique(AttackTechnique):
-    __metaclass__ = abc.ABCMeta
-
-    @staticmethod
-    def parse_usages(usage):
-        """
-        Parses data from database and translates usage enums into strings
-        :param usage: Usage telemetry that contains fields: {'usage': 'SMB', 'status': 1}
-        :return: usage string
-        """
-        try:
-            usage['usage'] = UsageEnum[usage['usage']].value[usage['status']]
-        except KeyError:
-            logger.error("Error translating usage enum. into string. "
-                         "Check if usage enum field exists and covers all telem. statuses.")
-        return usage
-
-    @classmethod
-    def get_usage_data(cls):
-        data = list(mongo.db.telemetry.aggregate(cls.get_usage_query()))
-        return list(map(cls.parse_usages, data))
-
-    @classmethod
-    def get_usage_query(cls):
-        """
-        :return: Query that parses attack telems for simple report component
-        (gets machines and attack technique usage).
-        """
-        return [{'$match': {'telem_category': 'attack',
-                            'data.technique': cls.tech_id}},
-                {'$lookup': {'from': 'monkey',
-                             'localField': 'monkey_guid',
-                             'foreignField': 'guid',
-                             'as': 'monkey'}},
-                {'$project': {'monkey': {'$arrayElemAt': ['$monkey', 0]},
-                              'status': '$data.status',
-                              'usage': '$data.usage'}},
-                {'$addFields': {'_id': 0,
-                                'machine': {'hostname': '$monkey.hostname', 'ips': '$monkey.ip_addresses'},
-                                'monkey': 0}},
-                {'$group': {'_id': {'machine': '$machine', 'status': '$status', 'usage': '$usage'}}},
-                {"$replaceRoot": {"newRoot": "$_id"}}]

monkey/monkey_island/cc/services/attack/technique_reports/usage_technique.py

@@ -0,0 +1,53 @@
+import abc
+
+from monkey_island.cc.database import mongo
+from monkey_island.cc.services.attack.technique_reports import AttackTechnique, logger
+from common.utils.attack_utils import UsageEnum
+
+
+class UsageTechnique(AttackTechnique):
+    __metaclass__ = abc.ABCMeta
+
+    @staticmethod
+    def parse_usages(usage):
+        """
+        Parses data from database and translates usage enums into strings
+        :param usage: Usage telemetry that contains fields: {'usage': 'SMB', 'status': 1}
+        :return: usage string
+        """
+        try:
+            usage['usage'] = UsageEnum[usage['usage']].value[usage['status']]
+        except KeyError:
+            logger.error("Error translating usage enum. into string. "
+                         "Check if usage enum field exists and covers all telem. statuses.")
+        return usage
+
+    @classmethod
+    def get_usage_data(cls):
+        """
+        Gets data of usage attack telemetries
+        :return: parsed list of usages from attack telemetries of usage type
+        """
+        data = list(mongo.db.telemetry.aggregate(cls.get_usage_query()))
+        return list(map(cls.parse_usages, data))
+
+    @classmethod
+    def get_usage_query(cls):
+        """
+        :return: Query that parses attack telemetries for a simple report component
+        (gets machines and attack technique usage).
+        """
+        return [{'$match': {'telem_category': 'attack',
+                            'data.technique': cls.tech_id}},
+                {'$lookup': {'from': 'monkey',
+                             'localField': 'monkey_guid',
+                             'foreignField': 'guid',
+                             'as': 'monkey'}},
+                {'$project': {'monkey': {'$arrayElemAt': ['$monkey', 0]},
+                              'status': '$data.status',
+                              'usage': '$data.usage'}},
+                {'$addFields': {'_id': 0,
+                                'machine': {'hostname': '$monkey.hostname', 'ips': '$monkey.ip_addresses'},
+                                'monkey': 0}},
+                {'$group': {'_id': {'machine': '$machine', 'status': '$status', 'usage': '$usage'}}},
+                {"$replaceRoot": {"newRoot": "$_id"}}]