From 7cb0e111ccd41c32a100e40d13de132d856fd293 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         | 26 +-----------
 .../attack/technique_reports/pba_technique.py | 41 +++++++++++++++++++
 2 files changed, 43 insertions(+), 24 deletions(-)
 create mode 100644 monkey/monkey_island/cc/services/attack/technique_reports/pba_technique.py

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 2db37bbbe..382fa6489 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/T1158.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1158.py
@@ -6,31 +6,9 @@ from monkey_island.cc.services.attack.technique_reports import AttackTechnique
 __author__ = "shreyamalviya"
 
 
-class T1158(AttackTechnique):
+class T1158(PostBreachTechnique):
     tech_id = "T1158"
     unscanned_msg = "Monkey did not 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."
-
-    query = [{'$match': {'telem_category': 'post_breach',
-                         'data.name': POST_BREACH_HIDDEN_FILES}},
-             {'$project': {'_id': 0,
-                           'machine': {'hostname': '$data.hostname',
-                                       'ips': ['$data.ip']},
-                           'result': '$data.result'}}]
-
-    @staticmethod
-    def get_report_data():
-        data = {'title': T1158.technique_title(), 'info': []}
-
-        hidden_file_info = list(mongo.db.telemetry.aggregate(T1158.query))
-
-        status = []
-        for pba_node in hidden_file_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(T1158.get_base_data_by_status(status))
-        data.update({'info': hidden_file_info})
-        return data
+    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
new file mode 100644
index 000000000..fce4edf70
--- /dev/null
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/pba_technique.py
@@ -0,0 +1,41 @@
+import abc
+
+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):
+    @property
+    @abc.abstractmethod
+    def pba_name(self):
+        """
+        :return: name of post breach action
+        """
+        pass
+
+    @classmethod
+    def get_pba_query(cls, post_breach_action_name):
+        return [{'$match': {'telem_category': 'post_breach',
+                            'data.name': post_breach_action_name}},
+                {'$project': {'_id': 0,
+                              'machine': {'hostname': '$data.hostname',
+                                          'ips': ['$data.ip']},
+                              'result': '$data.result'}}]
+
+    @classmethod
+    def get_report_data(cls):
+        data = {'title': cls.technique_title(), 'info': []}
+
+        info = list(mongo.db.telemetry.aggregate(cls.get_pba_query(cls.pba_name)))
+
+        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})
+        return data