Agent: Add PropagationCredentials type

This commit is contained in:
Mike Salvatore 2022-03-29 08:11:38 -04:00 committed by Ilija Lazoroski
parent e844ecf4e4
commit def62940af
2 changed files with 6 additions and 2 deletions

View File

@ -9,6 +9,7 @@ from typing import Callable, Dict, List, Mapping
from infection_monkey.i_puppet import ExploiterResultData, IPuppet
from infection_monkey.model import VictimHost
from infection_monkey.typing import PropagationCredentials
from infection_monkey.utils.threading import interruptible_iter, run_worker_threads
QUEUE_TIMEOUT = 2
@ -24,7 +25,7 @@ class Exploiter:
self,
puppet: IPuppet,
num_workers: int,
get_updated_credentials_for_propagation: Callable[[], Mapping],
get_updated_credentials_for_propagation: Callable[[], PropagationCredentials],
):
self._puppet = puppet
self._num_workers = num_workers
@ -160,7 +161,7 @@ class Exploiter:
exploitation_success=False, propagation_success=False, error_message=msg
)
def _get_credentials_for_propagation(self) -> Mapping:
def _get_credentials_for_propagation(self) -> PropagationCredentials:
try:
return self._get_updated_credentials_for_propagation()
except Exception as ex:

View File

@ -0,0 +1,3 @@
from typing import Iterable, Mapping
PropagationCredentials = Mapping[str, Iterable[str]]