From 9e6a569393e3e198d9a6579a78069a978689ab1a Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Mon, 15 Aug 2022 19:14:42 +0000
Subject: [PATCH] Agent: Update credentials repository to cache per-instance

---
 .../aggregating_propagation_credentials_repository.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py b/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py
index 3050e5f85..59ff279e0 100644
--- a/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py
+++ b/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Iterable, Sequence
+from typing import Any, Iterable
 
 from common.credentials import CredentialComponentType, Credentials, ICredentialComponent
 from infection_monkey.custom_types import PropagationCredentials
@@ -29,6 +29,11 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit
         }
         self._control_channel = control_channel
 
+        # Ensure caching happens per-instance instead of being shared across instances
+        self._get_credentials_from_control_channel = request_cache(CREDENTIALS_POLL_PERIOD_SEC)(
+            self._control_channel.get_credentials_for_propagation
+        )
+
     def add_credentials(self, credentials_to_add: Iterable[Credentials]):
         for credentials in credentials_to_add:
             if credentials.identity:
@@ -64,10 +69,6 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit
             self._stored_credentials = {}
             logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}")
 
-    @request_cache(CREDENTIALS_POLL_PERIOD_SEC)
-    def _get_credentials_from_control_channel(self) -> Sequence[Credentials]:
-        return self._control_channel.get_credentials_for_propagation()
-
     def _set_attribute(self, attribute_to_be_set: str, credentials_values: Iterable[Any]):
         if not credentials_values:
             return