From a099f21f61fd104ddaadc6194beb055a7688e5cb Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 13 Jun 2022 13:17:09 +0300 Subject: [PATCH] Agent: Initialize CustomPBA with a ControlClient object This is done to refactor ControlClient from a global --- monkey/infection_monkey/post_breach/custom_pba/custom_pba.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py index ce7dd64f2..f86df9e28 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -25,11 +25,12 @@ class CustomPBA(PBA): Defines user's configured post breach action. """ - def __init__(self, telemetry_messenger: ITelemetryMessenger): + def __init__(self, telemetry_messenger: ITelemetryMessenger, cc_client: ControlClient): super(CustomPBA, self).__init__( telemetry_messenger, POST_BREACH_FILE_EXECUTION, timeout=None ) self.filename = "" + self.cc_client = cc_client def run(self, options: Dict) -> Iterable[PostBreachData]: self._set_options(options) @@ -75,7 +76,7 @@ class CustomPBA(PBA): :return: True if successful, false otherwise """ - pba_file_contents = ControlClient.get_pba_file(filename) + pba_file_contents = self.cc_client.get_pba_file(filename) status = None if not pba_file_contents or not pba_file_contents.content: