From 1fdca52788d69ac865c51a022fc099b43fb56168 Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Wed, 22 Jul 2020 01:57:40 +0530
Subject: [PATCH] Fixes

---
 .../post_breach/actions/modify_shell_startup_files.py    | 4 ----
 .../cc/services/attack/technique_reports/T1156.py        | 9 +++++----
 .../cc/services/attack/technique_reports/T1504.py        | 9 +++++----
 .../attack/technique_reports/technique_report_tools.py   | 4 ++--
 4 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py b/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py
index 352610c6b..85f039628 100644
--- a/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py
+++ b/monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py
@@ -1,5 +1,3 @@
-import subprocess
-
 from common.data.post_breach_consts import \
     POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
 from infection_monkey.post_breach.pba import PBA
@@ -7,8 +5,6 @@ from infection_monkey.post_breach.shell_startup_files.shell_startup_files_modifi
     get_commands_to_modify_shell_startup_files
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 
-EXECUTION_WITHOUT_OUTPUT = "(PBA execution produced no output)"
-
 
 class ModifyShellStartupFiles(PBA):
     """
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 c6ddcba93..ed61bbc94 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/T1156.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1156.py
@@ -2,8 +2,9 @@ from common.data.post_breach_consts import \
     POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
 from monkey_island.cc.database import mongo
 from monkey_island.cc.services.attack.technique_reports import AttackTechnique
-from monkey_island.cc.services.attack.technique_reports.technique_report_tools import \
-    extract_shell_startup_files_modification_info, get_shell_startup_files_modification_status
+from monkey_island.cc.services.attack.technique_reports.technique_report_tools import (
+    extract_shell_startup_files_modification_info,
+    get_shell_startup_files_modification_status)
 
 __author__ = "shreyamalviya"
 
@@ -17,8 +18,8 @@ class T1156(AttackTechnique):
     query = [{'$match': {'telem_category': 'post_breach',
                          'data.name': POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION}},
              {'$project': {'_id': 0,
-                           'machine': {'hostname': '$data.hostname',
-                                       'ips': ['$data.ip']},
+                           'machine': {'hostname': {'$arrayElemAt': ['$data.hostname', 0]},
+                                       'ips': [{'$arrayElemAt': ['$data.ip', 0]}]},
                            'result': '$data.result'}}]
 
     @staticmethod
diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py
index 4fd0a475a..4da6ffd17 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1504.py
@@ -2,8 +2,9 @@ from common.data.post_breach_consts import \
     POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
 from monkey_island.cc.database import mongo
 from monkey_island.cc.services.attack.technique_reports import AttackTechnique
-from monkey_island.cc.services.attack.technique_reports.technique_report_tools import \
-    extract_shell_startup_files_modification_info, get_shell_startup_files_modification_status
+from monkey_island.cc.services.attack.technique_reports.technique_report_tools import (
+    extract_shell_startup_files_modification_info,
+    get_shell_startup_files_modification_status)
 
 __author__ = "shreyamalviya"
 
@@ -17,8 +18,8 @@ class T1504(AttackTechnique):
     query = [{'$match': {'telem_category': 'post_breach',
                          'data.name': POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION}},
              {'$project': {'_id': 0,
-                           'machine': {'hostname': '$data.hostname',
-                                       'ips': ['$data.ip']},
+                           'machine': {'hostname': {'$arrayElemAt': ['$data.hostname', 0]},
+                                       'ips': [{'$arrayElemAt': ['$data.ip', 0]}]},
                            'result': '$data.result'}}]
 
     @staticmethod
diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py b/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py
index e7e2eb8e0..88dbaab58 100644
--- a/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py
+++ b/monkey/monkey_island/cc/services/attack/technique_reports/technique_report_tools.py
@@ -1,5 +1,5 @@
-from monkey_island.cc.encryptor import encryptor
 from common.utils.attack_utils import ScanStatus
+from monkey_island.cc.encryptor import encryptor
 
 
 def parse_creds(attempt):
@@ -51,7 +51,7 @@ def extract_shell_startup_files_modification_info(shell_startup_files_modificati
     required_shell_startup_files_modification_info = []
     for shell_startup_file_result in shell_startup_files_modification_info[0]['result']:
         if any(file_name in shell_startup_file_result[0] for file_name in required_file_names):
-            shell_startup_files_modification_info.append({
+            required_shell_startup_files_modification_info.append({
                 'machine': shell_startup_files_modification_info[0]['machine'],
                 'result': shell_startup_file_result
                 })