From 8bb5096addf22e6aaf0d581f338c157fd981ffa5 Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Mon, 15 Jun 2020 17:45:23 +0530
Subject: [PATCH] Add T1504 report data + modify T1156 report data

---
 .../attack/technique_reports/T1156.py         | 19 +++++-----
 .../attack/technique_reports/T1504.py         | 38 +++++++++++++++++++
 2 files changed, 48 insertions(+), 9 deletions(-)
 create mode 100644 monkey/monkey_island/cc/services/attack/technique_reports/T1504.py

diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1156.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1156.py
index 96028865a..b618e743a 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/T1156.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1156.py
@@ -24,14 +24,15 @@ class T1156(AttackTechnique):
             if node['pba_results'] != 'None':
                 for pba in node['pba_results']:
                     if pba['name'] == POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION:
-                        status = ScanStatus.USED.value if pba['result'][1]\
-                                                       else ScanStatus.SCANNED.value
-                        data['info'].append({
-                                'machine': {
-                                    'hostname': pba['hostname'],
-                                    'ips': node['ip_addresses']
-                                },
-                                'result': pba['result'][0].replace('#', '')
-                            })
+                        if 'powershell.exe' not in pba['command']:
+                            status = ScanStatus.USED.value if pba['result'][1]\
+                                                        else ScanStatus.SCANNED.value
+                            data['info'].append({
+                                    'machine': {
+                                        'hostname': pba['hostname'],
+                                        'ips': node['ip_addresses']
+                                    },
+                                    'result': pba['result'][0].replace('#', '')
+                                })
             data.update(T1156.get_base_data_by_status(status))
         return data
diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py
new file mode 100644
index 000000000..585599c86
--- /dev/null
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py
@@ -0,0 +1,38 @@
+from monkey_island.cc.services.attack.technique_reports import AttackTechnique
+from monkey_island.cc.services.reporting.report import ReportService
+from common.utils.attack_utils import ScanStatus
+from common.data.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
+
+
+__author__ = "shreyamalviya"
+
+
+class T1504(AttackTechnique):
+    tech_id = "T1504"
+    unscanned_msg = "Monkey did not try modifying shell startup files on the system."
+    scanned_msg = "Monkey tried modifying shell startup files on the system but failed."
+    used_msg = "Monkey modified shell startup files on the system."
+
+    @staticmethod
+    def get_report_data():
+        data = {'title': T1504.technique_title(), 'info': []}
+
+        scanned_nodes = ReportService.get_scanned()
+        status = ScanStatus.UNSCANNED.value
+
+        for node in scanned_nodes:
+            if node['pba_results'] != 'None':
+                for pba in node['pba_results']:
+                    if pba['name'] == POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION:
+                        if 'powershell.exe' in pba['command']:
+                            status = ScanStatus.USED.value if pba['result'][1]\
+                                                        else ScanStatus.SCANNED.value
+                            data['info'].append({
+                                    'machine': {
+                                        'hostname': pba['hostname'],
+                                        'ips': node['ip_addresses']
+                                    },
+                                    'result': pba['result'][0].replace('#', '')
+                                })
+            data.update(T1504.get_base_data_by_status(status))
+        return data