Agent: Pass Agent ID to Heart

Issue #3119 PR #3155
guardicore · Mar 28, 2023 · a14a01c · a14a01c
1 parent 6cf118b
commit a14a01c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/monkey/infection_monkey/heart.py b/monkey/infection_monkey/heart.py
@@ -2,20 +2,20 @@
 import time
 
 from common.common_consts import HEARTBEAT_INTERVAL
+from common.types import AgentID
 from common.utils.code_utils import PeriodicCaller
 from infection_monkey.island_api_client import IIslandAPIClient
-from infection_monkey.utils.ids import get_agent_id
 
 logger = logging.getLogger(__name__)
 
 
 class Heart:
-    def __init__(self, island_api_client: IIslandAPIClient):
+    def __init__(self, island_api_client: IIslandAPIClient, agent_id: AgentID):
         self._island_api_client = island_api_client
         self._periodic_caller = PeriodicCaller(
             self._send_heartbeats, HEARTBEAT_INTERVAL, "AgentHeart"
         )
-        self._agent_id = get_agent_id()
+        self._agent_id = agent_id
 
     def start(self):
         logger.info("Starting the Agent's heart")

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
@@ -152,7 +152,7 @@ def __init__(self, args, ipc_logger_queue: multiprocessing.Queue, log_path: Path
             self._island_api_client, self._manager
         )
 
-        self._heart = Heart(self._island_api_client)
+        self._heart = Heart(self._island_api_client, self._agent_id)
         self._heart.start()
 
         self._current_depth = self._opts.depth