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

View File

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