Agent: Pass Agent ID to PluginRegistry

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

Hide file tree

Showing 4 changed files with 10 additions and 2 deletions.
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
@@ -402,6 +402,7 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet:
             self._propagation_credentials_repository,
             self._tcp_port_selector,
             otp_provider,
+            self._agent_id,
         )
         plugin_compatability_verifier = PluginCompatabilityVerifier(
             self._island_api_client, HARD_CODED_EXPLOITER_MANIFESTS

diff --git a/monkey/infection_monkey/puppet/plugin_registry.py b/monkey/infection_monkey/puppet/plugin_registry.py
@@ -9,12 +9,12 @@
 from common import OperatingSystem
 from common.agent_plugins import AgentPlugin, AgentPluginType
 from common.event_queue import IAgentEventPublisher
+from common.types import AgentID
 from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider
 from infection_monkey.i_puppet import UnknownPluginError
 from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIRequestError
 from infection_monkey.network import TCPPortSelector
 from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository
-from infection_monkey.utils.ids import get_agent_id
 
 from . import PluginSourceExtractor
 
@@ -35,6 +35,7 @@ def __init__(
         propagation_credentials_repository: IPropagationCredentialsRepository,
         tcp_port_selector: TCPPortSelector,
         otp_provider: IAgentOTPProvider,
+        agent_id: AgentID,
     ):
         """
         `self._registry` looks like -
@@ -57,7 +58,7 @@ def __init__(
         self._tcp_port_selector = tcp_port_selector
         self._otp_provider = otp_provider
 
-        self._agent_id = get_agent_id()
+        self._agent_id = agent_id
         self._lock = RLock()
 
     def get_plugin(self, plugin_type: AgentPluginType, plugin_name: str) -> Any:

diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py
@@ -6,6 +6,7 @@
 from common import OperatingSystem
 from common.agent_plugins import AgentPlugin, AgentPluginManifest, AgentPluginType
 from common.event_queue import IAgentEventPublisher
+from common.types import AgentID
 from infection_monkey.exploit import IAgentBinaryRepository, IAgentOTPProvider
 from infection_monkey.i_puppet import UnknownPluginError
 from infection_monkey.island_api_client import (
@@ -17,6 +18,8 @@
 from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository
 from infection_monkey.puppet import PluginRegistry, PluginSourceExtractor
 
+AGENT_ID = AgentID("707d801b-68cf-44d1-8a4e-7e1a89c412f8")
+
 
 @pytest.fixture
 def dummy_plugin_source_extractor() -> PluginSourceExtractor:
@@ -82,6 +85,7 @@ def test_get_plugin__error_handling(
         dummy_propagation_credentials_repository,
         dummy_tcp_port_selector,
         dummy_otp_provider,
+        AGENT_ID,
     )
 
     with pytest.raises(error_raised_by_plugin_registry):
@@ -147,6 +151,7 @@ def plugin_registry(
         dummy_propagation_credentials_repository,
         dummy_tcp_port_selector,
         dummy_otp_provider,
+        AGENT_ID,
     )
 
 

diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py
@@ -31,6 +31,7 @@ def mock_plugin_registry() -> PluginRegistry:
         MagicMock(),
         MagicMock(),
         MagicMock(),
+        AgentID("fd838244-385f-41b4-9904-495e3c7b644e"),
     )