From def62940af3df98ddcd506209e6522cdb4b59cc4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 29 Mar 2022 08:11:38 -0400 Subject: [PATCH] Agent: Add PropagationCredentials type --- monkey/infection_monkey/master/exploiter.py | 5 +++-- monkey/infection_monkey/typing.py | 3 +++ 2 files changed, 6 insertions(+), 2 deletions(-) create mode 100644 monkey/infection_monkey/typing.py diff --git a/monkey/infection_monkey/master/exploiter.py b/monkey/infection_monkey/master/exploiter.py index 3f0087af8..5f8e25b4d 100644 --- a/monkey/infection_monkey/master/exploiter.py +++ b/monkey/infection_monkey/master/exploiter.py @@ -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: diff --git a/monkey/infection_monkey/typing.py b/monkey/infection_monkey/typing.py new file mode 100644 index 000000000..6e1eab9aa --- /dev/null +++ b/monkey/infection_monkey/typing.py @@ -0,0 +1,3 @@ +from typing import Iterable, Mapping + +PropagationCredentials = Mapping[str, Iterable[str]]