From 936b9ead05addc6ebe618fe679ff2d2e4f748730 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Tue, 29 Mar 2022 10:18:37 +0300 Subject: [PATCH] Agent: Change post breach telem to use name from data argument --- .../infection_monkey/master/automated_master.py | 4 ++-- .../telemetry/post_breach_telem.py | 15 +++++---------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 9d37ceb8b..c1ced257c 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -198,8 +198,8 @@ class AutomatedMaster(IMaster): name = pba[0] options = pba[1] - display_name, result = self._puppet.run_pba(name, options) - self._telemetry_messenger.send_telemetry(PostBreachTelem(display_name, result)) + result = self._puppet.run_pba(name, options) + self._telemetry_messenger.send_telemetry(PostBreachTelem(result)) def _can_propagate(self) -> bool: return True diff --git a/monkey/infection_monkey/telemetry/post_breach_telem.py b/monkey/infection_monkey/telemetry/post_breach_telem.py index 19dddf76f..b968de71f 100644 --- a/monkey/infection_monkey/telemetry/post_breach_telem.py +++ b/monkey/infection_monkey/telemetry/post_breach_telem.py @@ -8,21 +8,16 @@ from infection_monkey.utils.environment import is_windows_os class PostBreachTelem(BaseTelem): - def __init__(self, name: str, post_breach_data: PostBreachData) -> None: - """ - Default post breach telemetry constructor - :param name: Name of post breach action - :param command: Command used as PBA - :param result: Result of PBA - """ + + telem_category = TelemCategoryEnum.POST_BREACH + + def __init__(self, post_breach_data: PostBreachData) -> None: super(PostBreachTelem, self).__init__() - self.name = name + self.name = post_breach_data.display_name self.command = post_breach_data.command self.result = post_breach_data.result self.hostname, self.ip = PostBreachTelem._get_hostname_and_ip() - telem_category = TelemCategoryEnum.POST_BREACH - def get_data(self) -> Dict: return { "command": self.command,