From def62940af3df98ddcd506209e6522cdb4b59cc4 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
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]]