From b237d258e52fbac1b52cd4f19baa99cba9526d68 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 13:39:33 +0530
Subject: [PATCH 01/10] Island: Use AgentConfiguration to get inaccessible
 subnets in ReportService

---
 monkey/monkey_island/cc/services/reporting/report.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py
index f57adc247..469f39ce0 100644
--- a/monkey/monkey_island/cc/services/reporting/report.py
+++ b/monkey/monkey_island/cc/services/reporting/report.py
@@ -11,9 +11,6 @@ from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.models.report import get_report, save_report
 from monkey_island.cc.repository import IAgentConfigurationRepository, ICredentialsRepository
-from monkey_island.cc.services.configuration.utils import (
-    get_config_network_segments_as_subnet_groups,
-)
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.reporting.exploitations.manual_exploitation import get_manual_monkeys
 from monkey_island.cc.services.reporting.exploitations.monkey_exploitation import (
@@ -325,8 +322,8 @@ class ReportService:
 
         return cross_segment_issues
 
-    @staticmethod
-    def get_cross_segment_issues():
+    @classmethod
+    def get_cross_segment_issues(cls):
         scans = mongo.db.telemetry.find(
             {"telem_category": "scan"},
             {
@@ -340,7 +337,8 @@ class ReportService:
         cross_segment_issues = []
 
         # For now the feature is limited to 1 group.
-        subnet_groups = get_config_network_segments_as_subnet_groups()
+        agent_configuration = cls._agent_configuration_repository.get_configuration()
+        subnet_groups = agent_configuration.propagation.network_scan.targets.inaccessible_subnets
 
         for subnet_group in subnet_groups:
             cross_segment_issues += ReportService.get_cross_segment_issues_per_subnet_group(

From 093e6fa3a6666214fe1f9c8b245738e80504c3d4 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:14:33 +0530
Subject: [PATCH 02/10] Island: Add init and IAgentConfigurationRepository
 dependency to Telemetry resource

---
 monkey/monkey_island/cc/resources/telemetry.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py
index 8b6bf790a..32e8e1a23 100644
--- a/monkey/monkey_island/cc/resources/telemetry.py
+++ b/monkey/monkey_island/cc/resources/telemetry.py
@@ -8,6 +8,7 @@ from flask import request
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models.monkey import Monkey
 from monkey_island.cc.models.telemetries import get_telemetry_by_query
+from monkey_island.cc.repository import IAgentConfigurationRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
 from monkey_island.cc.services.node import NodeService
@@ -20,6 +21,9 @@ class Telemetry(AbstractResource):
     # API Spec: Resource name should be plural
     urls = ["/api/telemetry", "/api/telemetry/<string:monkey_guid>"]
 
+    def __init__(self, agent_configuration_repository: IAgentConfigurationRepository):
+        self._agent_configuration_repository = agent_configuration_repository
+
     @jwt_required
     def get(self, **kw):
         monkey_guid = request.args.get("monkey_guid")

From c70627c74ba3d0979d180ac9dd06efcb985eee77 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:16:19 +0530
Subject: [PATCH 03/10] Island: Pass agent configuration to telemetry
 processing function in Telemetry resource

---
 monkey/monkey_island/cc/resources/telemetry.py                 | 3 ++-
 .../cc/services/telemetry/processing/processing.py             | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py
index 32e8e1a23..106077428 100644
--- a/monkey/monkey_island/cc/resources/telemetry.py
+++ b/monkey/monkey_island/cc/resources/telemetry.py
@@ -63,7 +63,8 @@ class Telemetry(AbstractResource):
         monkey = NodeService.get_monkey_by_guid(telemetry_json["monkey_guid"])
         NodeService.update_monkey_modify_time(monkey["_id"])
 
-        process_telemetry(telemetry_json)
+        agent_configuration = self._agent_configuration_repository.get_configuration()
+        process_telemetry(telemetry_json, agent_configuration)
 
         # API Spec: RESTful way is to return an identifier of the updated/newly created resource
         return {}, 201
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/processing.py b/monkey/monkey_island/cc/services/telemetry/processing/processing.py
index 569957c75..a1db59b45 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/processing.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/processing.py
@@ -1,6 +1,7 @@
 import logging
 
 from common.common_consts.telem_categories import TelemCategoryEnum
+from common.configuration import AgentConfiguration
 from monkey_island.cc.models.telemetries import save_telemetry
 from monkey_island.cc.services.telemetry.processing.aws_info import process_aws_telemetry
 from monkey_island.cc.services.telemetry.processing.exploit import process_exploit_telemetry
@@ -29,7 +30,7 @@ TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = {
 UNSAVED_TELEMETRIES = [TelemCategoryEnum.CREDENTIALS]
 
 
-def process_telemetry(telemetry_json):
+def process_telemetry(telemetry_json, agent_configuration: AgentConfiguration):
     try:
         telem_category = telemetry_json.get("telem_category")
         if telem_category in TELEMETRY_CATEGORY_TO_PROCESSING_FUNC:

From 8daa6db81f0233c9ebc6c8db8ed1e58d6366316d Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:18:38 +0530
Subject: [PATCH 04/10] Island: Modify all telemetry processing functions to
 accept an extra agent configuration argument

---
 .../cc/services/telemetry/processing/aws_info.py              | 2 +-
 .../monkey_island/cc/services/telemetry/processing/exploit.py | 2 +-
 .../cc/services/telemetry/processing/post_breach.py           | 2 +-
 .../cc/services/telemetry/processing/processing.py            | 4 +++-
 monkey/monkey_island/cc/services/telemetry/processing/scan.py | 2 +-
 .../monkey_island/cc/services/telemetry/processing/state.py   | 3 ++-
 .../monkey_island/cc/services/telemetry/processing/tunnel.py  | 2 +-
 7 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/monkey/monkey_island/cc/services/telemetry/processing/aws_info.py b/monkey/monkey_island/cc/services/telemetry/processing/aws_info.py
index 020f236f0..1821b9e15 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/aws_info.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/aws_info.py
@@ -5,7 +5,7 @@ from monkey_island.cc.models.monkey import Monkey
 logger = logging.getLogger(__name__)
 
 
-def process_aws_telemetry(telemetry_json):
+def process_aws_telemetry(telemetry_json, _):
     relevant_monkey = Monkey.get_single_monkey_by_guid(telemetry_json["monkey_guid"])
 
     if "instance_id" in telemetry_json["data"]:
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/exploit.py b/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
index dc5b2e638..cc6bd2c03 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/exploit.py
@@ -14,7 +14,7 @@ from monkey_island.cc.services.telemetry.zero_trust_checks.machine_exploited imp
 )
 
 
-def process_exploit_telemetry(telemetry_json):
+def process_exploit_telemetry(telemetry_json, _):
     encrypt_exploit_creds(telemetry_json)
     edge = get_edge_by_scan_or_exploit_telemetry(telemetry_json)
     update_network_with_exploit(edge, telemetry_json)
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 e4f83947e..401589027 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/post_breach.py
@@ -32,7 +32,7 @@ POST_BREACH_TELEMETRY_PROCESSING_FUNCS = {
 }
 
 
-def process_post_breach_telemetry(telemetry_json):
+def process_post_breach_telemetry(telemetry_json, _):
     def convert_telem_data_to_list(data):
         modified_data = [data]
         if type(data["result"][0]) is list:  # multiple results in one pba
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/processing.py b/monkey/monkey_island/cc/services/telemetry/processing/processing.py
index a1db59b45..ee2d2d065 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/processing.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/processing.py
@@ -34,7 +34,9 @@ def process_telemetry(telemetry_json, agent_configuration: AgentConfiguration):
     try:
         telem_category = telemetry_json.get("telem_category")
         if telem_category in TELEMETRY_CATEGORY_TO_PROCESSING_FUNC:
-            TELEMETRY_CATEGORY_TO_PROCESSING_FUNC[telem_category](telemetry_json)
+            TELEMETRY_CATEGORY_TO_PROCESSING_FUNC[telem_category](
+                telemetry_json, agent_configuration
+            )
         else:
             logger.info("Got unknown type of telemetry: %s" % telem_category)
 
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scan.py b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
index 54379dc45..2a29b7762 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/scan.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
@@ -14,7 +14,7 @@ from monkey_island.cc.services.telemetry.zero_trust_checks.segmentation import (
 )
 
 
-def process_scan_telemetry(telemetry_json):
+def process_scan_telemetry(telemetry_json, _):
     if not _host_responded(telemetry_json["data"]["machine"]):
         return
 
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/state.py b/monkey/monkey_island/cc/services/telemetry/processing/state.py
index 87e7797c2..f0146fce7 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/state.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/state.py
@@ -1,5 +1,6 @@
 import logging
 
+from common.configuration import AgentConfiguration
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.telemetry.zero_trust_checks.segmentation import (
@@ -9,7 +10,7 @@ from monkey_island.cc.services.telemetry.zero_trust_checks.segmentation import (
 logger = logging.getLogger(__name__)
 
 
-def process_state_telemetry(telemetry_json):
+def process_state_telemetry(telemetry_json, agent_configuration: AgentConfiguration):
     monkey = NodeService.get_monkey_by_guid(telemetry_json["monkey_guid"])
     NodeService.add_communication_info(monkey, telemetry_json["command_control_channel"])
     if telemetry_json["data"]["done"]:
diff --git a/monkey/monkey_island/cc/services/telemetry/processing/tunnel.py b/monkey/monkey_island/cc/services/telemetry/processing/tunnel.py
index 4464eb82a..6bd1fd711 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/tunnel.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/tunnel.py
@@ -5,7 +5,7 @@ from monkey_island.cc.services.telemetry.zero_trust_checks.tunneling import (
 )
 
 
-def process_tunnel_telemetry(telemetry_json):
+def process_tunnel_telemetry(telemetry_json, _):
     check_tunneling_violation(telemetry_json)
     monkey_id = NodeService.get_monkey_by_guid(telemetry_json["monkey_guid"])["_id"]
     if telemetry_json["data"]["proxy"] is not None:

From df6296fc6befa8387aeb01a7364a60517d62148b Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:20:27 +0530
Subject: [PATCH 05/10] Island: Use agent configuration in ZT segmentation
 check to get inaccessible subnets

---
 .../cc/services/telemetry/processing/state.py       |  2 +-
 .../telemetry/zero_trust_checks/segmentation.py     | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/monkey/monkey_island/cc/services/telemetry/processing/state.py b/monkey/monkey_island/cc/services/telemetry/processing/state.py
index f0146fce7..208fe0c72 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/state.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/state.py
@@ -20,7 +20,7 @@ def process_state_telemetry(telemetry_json, agent_configuration: AgentConfigurat
 
     if telemetry_json["data"]["done"]:
         current_monkey = Monkey.get_single_monkey_by_guid(telemetry_json["monkey_guid"])
-        check_passed_findings_for_unreached_segments(current_monkey)
+        check_passed_findings_for_unreached_segments(current_monkey, agent_configuration)
 
     if telemetry_json["data"]["version"]:
         logger.info(
diff --git a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py
index d26e2bd69..47a24d468 100644
--- a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py
+++ b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py
@@ -1,6 +1,7 @@
 import itertools
 
 import common.common_consts.zero_trust_consts as zero_trust_consts
+from common.configuration import AgentConfiguration
 from common.network.network_range import NetworkRange
 from common.network.segmentation_utils import get_ip_if_in_subnet, get_ip_in_src_and_not_in_dst
 from monkey_island.cc.models import Monkey
@@ -84,13 +85,21 @@ def get_segmentation_violation_event(current_monkey, source_subnet, target_ip, t
     )
 
 
-def check_passed_findings_for_unreached_segments(current_monkey):
+def check_passed_findings_for_unreached_segments(
+    current_monkey, agent_configuration: AgentConfiguration
+):
     flat_all_subnets = [
-        item for sublist in get_config_network_segments_as_subnet_groups() for item in sublist
+        item
+        for sublist in _get_config_network_segments_as_subnet_groups(agent_configuration)
+        for item in sublist
     ]
     create_or_add_findings_for_all_pairs(flat_all_subnets, current_monkey)
 
 
+def _get_config_network_segments_as_subnet_groups(agent_configuration: AgentConfiguration):
+    return agent_configuration.propagation.network_scan.targets.inaccessible_subnets
+
+
 def create_or_add_findings_for_all_pairs(all_subnets, current_monkey):
     # Filter the subnets that this monkey is part of.
     this_monkey_subnets = []

From 45151e7e19e1d256078a3da9f0c3be8f484ffee9 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:24:10 +0530
Subject: [PATCH 06/10] Island: Modify scan telemetry processing function to
 accept agent configuration

---
 monkey/monkey_island/cc/services/telemetry/processing/scan.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scan.py b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
index 2a29b7762..7c31c7761 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/scan.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
@@ -1,5 +1,6 @@
 from typing import Mapping
 
+from common.configuration import AgentConfiguration
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.services.node import NodeService
@@ -14,7 +15,7 @@ from monkey_island.cc.services.telemetry.zero_trust_checks.segmentation import (
 )
 
 
-def process_scan_telemetry(telemetry_json, _):
+def process_scan_telemetry(telemetry_json, agent_configuration: AgentConfiguration):
     if not _host_responded(telemetry_json["data"]["machine"]):
         return
 

From d534c64a3fcec310b1f806cfcaefeaa4b91f45b5 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:25:01 +0530
Subject: [PATCH 07/10] Island: Use agent configuration in remaining ZT
 segmentation check function to get inaccessible subnets

---
 .../cc/services/telemetry/processing/scan.py             | 2 +-
 .../services/telemetry/zero_trust_checks/segmentation.py | 9 ++++-----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scan.py b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
index 7c31c7761..52a9325e4 100644
--- a/monkey/monkey_island/cc/services/telemetry/processing/scan.py
+++ b/monkey/monkey_island/cc/services/telemetry/processing/scan.py
@@ -24,7 +24,7 @@ def process_scan_telemetry(telemetry_json, agent_configuration: AgentConfigurati
 
     current_monkey = Monkey.get_single_monkey_by_guid(telemetry_json["monkey_guid"])
     target_ip = telemetry_json["data"]["machine"]["ip_addr"]
-    check_segmentation_violation(current_monkey, target_ip)
+    check_segmentation_violation(current_monkey, target_ip, agent_configuration)
 
 
 def update_edges_and_nodes_based_on_scan_telemetry(telemetry_json):
diff --git a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py
index 47a24d468..5411e7600 100644
--- a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py
+++ b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py
@@ -6,9 +6,6 @@ from common.network.network_range import NetworkRange
 from common.network.segmentation_utils import get_ip_if_in_subnet, get_ip_in_src_and_not_in_dst
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.models.zero_trust.event import Event
-from monkey_island.cc.services.configuration.utils import (
-    get_config_network_segments_as_subnet_groups,
-)
 from monkey_island.cc.services.zero_trust.monkey_findings.monkey_zt_finding_service import (
     MonkeyZTFindingService,
 )
@@ -25,9 +22,11 @@ SEGMENTATION_VIOLATION_EVENT_TEXT = (
 )
 
 
-def check_segmentation_violation(current_monkey, target_ip):
+def check_segmentation_violation(
+    current_monkey, target_ip, agent_configuration: AgentConfiguration
+):
     # TODO - lower code duplication between this and report.py.
-    subnet_groups = get_config_network_segments_as_subnet_groups()
+    subnet_groups = _get_config_network_segments_as_subnet_groups(agent_configuration)
     for subnet_group in subnet_groups:
         subnet_pairs = itertools.product(subnet_group, subnet_group)
         for subnet_pair in subnet_pairs:

From 02a06787d0bad926666147b09952f316bca8e57d Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:26:30 +0530
Subject: [PATCH 08/10] Island: Remove
 monkey/monkey_island/cc/services/configuration/

---
 monkey/monkey_island/cc/services/configuration/__init__.py | 0
 monkey/monkey_island/cc/services/configuration/utils.py    | 6 ------
 2 files changed, 6 deletions(-)
 delete mode 100644 monkey/monkey_island/cc/services/configuration/__init__.py
 delete mode 100644 monkey/monkey_island/cc/services/configuration/utils.py

diff --git a/monkey/monkey_island/cc/services/configuration/__init__.py b/monkey/monkey_island/cc/services/configuration/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/monkey/monkey_island/cc/services/configuration/utils.py b/monkey/monkey_island/cc/services/configuration/utils.py
deleted file mode 100644
index 8d2b5c18c..000000000
--- a/monkey/monkey_island/cc/services/configuration/utils.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from common.config_value_paths import INACCESSIBLE_SUBNETS_PATH
-from monkey_island.cc.services.config import ConfigService
-
-
-def get_config_network_segments_as_subnet_groups():
-    return [ConfigService.get_config_value(INACCESSIBLE_SUBNETS_PATH)]

From d6bd84ad66a979a375b6dc07fb4d00f602a1126b Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:28:49 +0530
Subject: [PATCH 09/10] Common: Remove unused config value path
 INACCESSIBLE_SUBNETS_PATH

---
 monkey/common/config_value_paths.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/monkey/common/config_value_paths.py b/monkey/common/config_value_paths.py
index 77fc0fb60..5e29ebf11 100644
--- a/monkey/common/config_value_paths.py
+++ b/monkey/common/config_value_paths.py
@@ -1,5 +1,4 @@
 SSH_KEYS_PATH = ["internal", "exploits", "exploit_ssh_keys"]
-INACCESSIBLE_SUBNETS_PATH = ["basic_network", "network_analysis", "inaccessible_subnets"]
 USER_LIST_PATH = ["basic", "credentials", "exploit_user_list"]
 PASSWORD_LIST_PATH = ["basic", "credentials", "exploit_password_list"]
 LM_HASH_LIST_PATH = ["internal", "exploits", "exploit_lm_hash_list"]

From 68a7085ede233302ae899d3d5fc4068cc397576e Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Thu, 14 Jul 2022 14:49:11 +0530
Subject: [PATCH 10/10] UT: Add placeholder argument to function call in
 test_process_post_breach_telemetry

---
 .../cc/services/telemetry/processing/test_post_breach.py      | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/telemetry/processing/test_post_breach.py b/monkey/tests/unit_tests/monkey_island/cc/services/telemetry/processing/test_post_breach.py
index f6d33b930..01a847258 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/services/telemetry/processing/test_post_breach.py
+++ b/monkey/tests/unit_tests/monkey_island/cc/services/telemetry/processing/test_post_breach.py
@@ -69,8 +69,8 @@ expected_telem_single_result = {
 def test_process_post_breach_telemetry():
     post_breach.update_data = Mock()  # actual behavior of update_data() is to access mongodb
     # multiple results in PBA
-    post_breach.process_post_breach_telemetry(original_telem_multiple_results)
+    post_breach.process_post_breach_telemetry(original_telem_multiple_results, None)
     assert original_telem_multiple_results == expected_telem_multiple_results
     # single result in PBA
-    post_breach.process_post_breach_telemetry(original_telem_single_result)
+    post_breach.process_post_breach_telemetry(original_telem_single_result, None)
     assert original_telem_single_result == expected_telem_single_result