From 3c853b662541fd38320b8691098a223049cc676c Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 28 Mar 2022 14:32:14 +0300 Subject: [PATCH] Agent: Change PostBreachTelemetry to accept post breach data --- monkey/infection_monkey/master/automated_master.py | 4 ++-- monkey/infection_monkey/telemetry/post_breach_telem.py | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 251240947..9d37ceb8b 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, command, result = self._puppet.run_pba(name, options) - self._telemetry_messenger.send_telemetry(PostBreachTelem(display_name, command, result)) + display_name, result = self._puppet.run_pba(name, options) + self._telemetry_messenger.send_telemetry(PostBreachTelem(display_name, 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 e4f93e30d..19dddf76f 100644 --- a/monkey/infection_monkey/telemetry/post_breach_telem.py +++ b/monkey/infection_monkey/telemetry/post_breach_telem.py @@ -2,12 +2,13 @@ import socket from typing import Dict, Tuple from common.common_consts.telem_categories import TelemCategoryEnum +from infection_monkey.i_puppet import PostBreachData from infection_monkey.telemetry.base_telem import BaseTelem from infection_monkey.utils.environment import is_windows_os class PostBreachTelem(BaseTelem): - def __init__(self, name: str, command: str, result: str) -> None: + def __init__(self, name: str, post_breach_data: PostBreachData) -> None: """ Default post breach telemetry constructor :param name: Name of post breach action @@ -16,8 +17,8 @@ class PostBreachTelem(BaseTelem): """ super(PostBreachTelem, self).__init__() self.name = name - self.command = command - self.result = result + 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