p34709852
/
monkey
forked from p15670423/monkey
1
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Initialize CustomPBA with a ControlClient object 

This is done to refactor ControlClient from a global
This commit is contained in:
vakarisz 2022-06-13 13:17:09 +03:00
parent 049eb1b174
commit a099f21f61
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
monkey/infection_monkey/post_breach/custom_pba/custom_pba.py
View File

@ -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: