From d25ad3a209bd1d9e1c7b5e4e312b59e61362ee8c Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Fri, 10 Jul 2020 20:42:37 +0530
Subject: [PATCH] Refactor PBA report data code

Only with T1158 for now
---
 .../attack/technique_reports/T1158.py         |  2 +-
 .../attack/technique_reports/pba_technique.py | 35 ++++++-------------
 2 files changed, 12 insertions(+), 25 deletions(-)

diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1158.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1158.py
index 7b0f87358..dce2b3adb 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/T1158.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1158.py
@@ -10,4 +10,4 @@ class T1158(PostBreachTechnique):
     unscanned_msg = "Monkey didn't try creating hidden files or folders."
     scanned_msg = "Monkey tried creating hidden files and folders on the system but failed."
     used_msg = "Monkey created hidden files and folders on the system."
-    pba_names = [POST_BREACH_HIDDEN_FILES]
+    pba_name = POST_BREACH_HIDDEN_FILES
diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/pba_technique.py b/monkey/monkey_island/cc/services/attack/technique_reports/pba_technique.py
index a7ef96803..fce4edf70 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/pba_technique.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/pba_technique.py
@@ -1,32 +1,24 @@
 import abc
 
-from common.utils.attack_utils import ScanStatus
-from monkey_island.cc.database import mongo
 from monkey_island.cc.services.attack.attack_config import AttackConfig
+from monkey_island.cc.database import mongo
+from common.utils.attack_utils import ScanStatus
 from monkey_island.cc.services.attack.technique_reports import AttackTechnique
 
 
 class PostBreachTechnique(AttackTechnique, metaclass=abc.ABCMeta):
-    """ Class for ATT&CK report components of post-breach actions """
-
     @property
     @abc.abstractmethod
-    def pba_names(self):
+    def pba_name(self):
         """
         :return: name of post breach action
         """
         pass
 
     @classmethod
-    def get_pba_query(cls, post_breach_action_names):
-        """
-        :param post_breach_action_names: Names of post-breach actions with which the technique is associated
-        (example - `["Communicate as new user", "Backdoor user"]` for T1136)
-        :return: Mongo query that parses attack telemetries for a simple report component
-        (gets machines and post-breach action usage).
-        """
+    def get_pba_query(cls, post_breach_action_name):
         return [{'$match': {'telem_category': 'post_breach',
-                            '$or': [{'data.name': pba_name} for pba_name in post_breach_action_names]}},
+                            'data.name': post_breach_action_name}},
                 {'$project': {'_id': 0,
                               'machine': {'hostname': '$data.hostname',
                                           'ips': ['$data.ip']},
@@ -34,20 +26,15 @@ class PostBreachTechnique(AttackTechnique, metaclass=abc.ABCMeta):
 
     @classmethod
     def get_report_data(cls):
-        """
-        :return: Technique's report data aggregated from the database
-        """
         data = {'title': cls.technique_title(), 'info': []}
 
-        info = list(mongo.db.telemetry.aggregate(cls.get_pba_query(cls.pba_names)))
+        info = list(mongo.db.telemetry.aggregate(cls.get_pba_query(cls.pba_name)))
 
-        status = ScanStatus.UNSCANNED.value
-        if info:
-            successful_PBAs = mongo.db.telemetry.count({
-                '$or': [{'data.name': pba_name} for pba_name in cls.pba_names],
-                'data.result.1': True
-            })
-            status = ScanStatus.USED.value if successful_PBAs else ScanStatus.SCANNED.value
+        status = []
+        for pba_node in info:
+            status.append(pba_node['result'][1])
+        status = (ScanStatus.USED.value if any(status) else ScanStatus.SCANNED.value)\
+            if status else ScanStatus.UNSCANNED.value
 
         data.update(cls.get_base_data_by_status(status))
         data.update({'info': info})