Agent: Fix typehints in _run_pbas of automated_master.py

Typehint was Mapping, when it was using and calling other methods with CustomPBAConfiguration
This commit is contained in:
vakarisz 2022-06-27 09:36:56 +03:00
parent fefd2daf2b
commit 4b7ab058c6
1 changed files with 4 additions and 4 deletions

View File

@ -1,9 +1,9 @@
import logging import logging
import threading import threading
import time 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 common.utils import Timer
from infection_monkey.credential_store import ICredentialsStore from infection_monkey.credential_store import ICredentialsStore
from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
@ -208,13 +208,13 @@ class AutomatedMaster(IMaster):
self, self,
plugins: Iterable[PluginConfiguration], plugins: Iterable[PluginConfiguration],
callback: Callable[[Any], None], callback: Callable[[Any], None],
custom_pba_options: Mapping, custom_pba_options: CustomPBAConfiguration,
): ):
self._run_plugins(plugins, "post-breach action", callback) self._run_plugins(plugins, "post-breach action", callback)
if custom_pba_is_enabled(custom_pba_options): if custom_pba_is_enabled(custom_pba_options):
self._run_plugins( self._run_plugins(
[PluginConfiguration(name="CustomPBA", options=custom_pba_options)], [PluginConfiguration(name="CustomPBA", options=custom_pba_options.__dict__)],
"post-breach action", "post-breach action",
callback, callback,
) )