From 4b7ab058c6b69e941c44bb4881f809cc57f577be Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 27 Jun 2022 09:36:56 +0300 Subject: [PATCH] Agent: Fix typehints in _run_pbas of automated_master.py Typehint was Mapping, when it was using and calling other methods with CustomPBAConfiguration --- monkey/infection_monkey/master/automated_master.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 9333554b7..ad73b50f4 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -1,9 +1,9 @@ import logging import threading import time -from typing import Any, Callable, Iterable, List, Mapping, Optional +from typing import Any, Callable, Iterable, List, Optional -from common.configuration import PluginConfiguration +from common.configuration import CustomPBAConfiguration, PluginConfiguration from common.utils import Timer from infection_monkey.credential_store import ICredentialsStore from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError @@ -208,13 +208,13 @@ class AutomatedMaster(IMaster): self, plugins: Iterable[PluginConfiguration], callback: Callable[[Any], None], - custom_pba_options: Mapping, + custom_pba_options: CustomPBAConfiguration, ): self._run_plugins(plugins, "post-breach action", callback) if custom_pba_is_enabled(custom_pba_options): self._run_plugins( - [PluginConfiguration(name="CustomPBA", options=custom_pba_options)], + [PluginConfiguration(name="CustomPBA", options=custom_pba_options.__dict__)], "post-breach action", callback, )