Agent: Move the decision if custom pba should run to master

This commit is contained in:
vakarisz 2022-03-30 16:37:19 +03:00
parent e855d2ed34
commit 29a545a58f
3 changed files with 15 additions and 12 deletions

View File

@ -9,7 +9,6 @@ from infection_monkey.i_master import IMaster
from infection_monkey.i_puppet import IPuppet
from infection_monkey.model import VictimHostFactory
from infection_monkey.network import NetworkInterface
from infection_monkey.post_breach.custom_pba import CustomPBA
from infection_monkey.telemetry.credentials_telem import CredentialsTelem
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
@ -17,6 +16,7 @@ from infection_monkey.utils.threading import create_daemon_thread, interruptible
from infection_monkey.utils.timer import Timer
from . import Exploiter, IPScanner, Propagator
from .option_parsing import custom_pba_is_enabled
CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5
CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5
@ -214,7 +214,7 @@ class AutomatedMaster(IMaster):
):
self._run_plugins(plugins, "post-breach action", callback)
if CustomPBA.should_run(custom_pba_options):
if custom_pba_is_enabled(custom_pba_options):
self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback)
def _run_plugins(

View File

@ -0,0 +1,13 @@
from typing import Dict
from infection_monkey.utils.environment import is_windows_os
def custom_pba_is_enabled(pba_options: Dict) -> bool:
if not is_windows_os():
if pba_options["linux_command"]:
return True
else:
if pba_options["windows_command"]:
return True
return False

View File

@ -65,16 +65,6 @@ class CustomPBA(PBA):
self.download_pba_file(get_monkey_dir_path(), self.filename)
return super(CustomPBA, self)._execute_default()
@staticmethod
def should_run(options):
if not is_windows_os():
if options["linux_filename"] or options["linux_command"]:
return True
else:
if options["windows_filename"] or options["windows_command"]:
return True
return False
def download_pba_file(self, dst_dir, filename):
"""
Handles post breach action file download