From e674f9e0c021e82d7beed2251b45024f3c18d17d Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Wed, 16 Feb 2022 16:41:00 +0530
Subject: [PATCH] Island: Move antivirus check for ZT report from system info
 processing to PBA processing

---
 .../telemetry/processing/post_breach.py         | 17 ++++++++++++++++-
 .../system_info_telemetry_dispatcher.py         |  4 +---
 .../zero_trust_checks/antivirus_existence.py    |  7 ++-----
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py b/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
index 5506ff54d..3e02971de 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
@@ -1,8 +1,14 @@
 import copy
 
-from common.common_consts.post_breach_consts import POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER
+from common.common_consts.post_breach_consts import (
+    POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER,
+    POST_BREACH_PROCESS_LIST_COLLECTION,
+)
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
+from monkey_island.cc.services.telemetry.zero_trust_checks.antivirus_existence import (
+    check_antivirus_existence,
+)
 from monkey_island.cc.services.telemetry.zero_trust_checks.communicate_as_backdoor_user import (
     check_new_user_communication,
 )
@@ -17,8 +23,17 @@ def process_communicate_as_backdoor_user_telemetry(telemetry_json):
     check_new_user_communication(current_monkey, success, message)
 
 
+def process_process_list_collection_telemetry(telemetry_json):
+    current_monkey = Monkey.get_single_monkey_by_guid(telemetry_json["monkey_guid"])
+    check_antivirus_existence(telemetry_json, current_monkey)
+
+
 POST_BREACH_TELEMETRY_PROCESSING_FUNCS = {
     POST_BREACH_COMMUNICATE_AS_BACKDOOR_USER: process_communicate_as_backdoor_user_telemetry,
+    # TODO: Remove line 31 and un-comment line 32 after the TODO in `_run_pba()` in
+    #       `automated_master.py` is resolved.
+    "ProcessListCollection": process_process_list_collection_telemetry,
+    # POST_BREACH_PROCESS_LIST_COLLECTION: process_process_list_collection_telemetry,
 }
 
 
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py
index 9df25a677..7faae8eb2 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py
@@ -1,12 +1,10 @@
 import logging
 import typing
 
-from monkey_island.cc.services.telemetry.zero_trust_checks.antivirus_existence import (
-    check_antivirus_existence,
-)
 
 logger = logging.getLogger(__name__)
 
+
 SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSORS = {}
 
 
diff --git a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/antivirus_existence.py b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/antivirus_existence.py
index d2f154a9e..4e8a86fb4 100644
--- a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/antivirus_existence.py
+++ b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/antivirus_existence.py
@@ -1,7 +1,6 @@
 import json
 
 import common.common_consts.zero_trust_consts as zero_trust_consts
-from monkey_island.cc.models import Monkey
 from monkey_island.cc.models.zero_trust.event import Event
 from monkey_island.cc.services.telemetry.zero_trust_checks.known_anti_viruses import (
     ANTI_VIRUS_KNOWN_PROCESS_NAMES,
@@ -11,9 +10,7 @@ from monkey_island.cc.services.zero_trust.monkey_findings.monkey_zt_finding_serv
 )
 
 
-def check_antivirus_existence(process_list_json, monkey_guid):
-    current_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
-
+def check_antivirus_existence(telemetry_json, current_monkey):
     process_list_event = Event.create_event(
         title="Process list",
         message="Monkey on {} scanned the process list".format(current_monkey.hostname),
@@ -21,7 +18,7 @@ def check_antivirus_existence(process_list_json, monkey_guid):
     )
     events = [process_list_event]
 
-    av_processes = filter_av_processes(process_list_json["process_list"])
+    av_processes = filter_av_processes(telemetry_json["data"]["result"][0])
 
     for process in av_processes:
         events.append(