UT: Fix tests to work with new Agent ID logic

monkey/tests/unit_tests/agent_plugins/exploiters/hadoop/test_hadoop_exploiter.py

@@ -18,7 +18,9 @@
     TargetHost,
     TargetHostPorts,
 )
+from infection_monkey.utils.ids import get_agent_id
 
+AGENT_ID = get_agent_id()
 TARGET_IP = IPv4Address("1.1.1.1")
 SERVERS = ["10.10.10.10"]
 DOWNLOAD_URL = "http://download.me"
@@ -83,7 +85,7 @@ def hadoop_exploiter(
     mock_hadoop_exploit_client: HadoopExploitClient,
     mock_start_agent_binary_server: Callable[[TargetHost], HTTPBytesServer],
 ) -> HadoopExploiter:
-    return HadoopExploiter(mock_hadoop_exploit_client, mock_start_agent_binary_server)
+    return HadoopExploiter(AGENT_ID, mock_hadoop_exploit_client, mock_start_agent_binary_server)
 
 
 @pytest.fixture

monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_command_builder.py

@@ -7,10 +7,12 @@
 from common import OperatingSystem
 from infection_monkey.exploit import IAgentOTPProvider
 from infection_monkey.model import DROPPER_ARG, MONKEY_ARG
+from infection_monkey.utils.ids import get_agent_id
 
 DROPPER_EXE_PATH = PureWindowsPath("C:\\dropper.exe")
 AGENT_EXE_PATH = PureWindowsPath("C:\\agent.exe")
 OTP = "123456"
+AGENT_ID = get_agent_id()
 
 
 @pytest.fixture
@@ -23,6 +25,7 @@ def otp_provider() -> IAgentOTPProvider:
 def test_exception_raised_for_linux(otp_provider: IAgentOTPProvider):
     with pytest.raises(Exception):
         build_smb_command(
+            AGENT_ID,
             ["127.0.0.1"],
             2,
             OperatingSystem.LINUX,
@@ -46,7 +49,13 @@ def test_servers(
 ):
     servers = ["127.0.0.1", "192.168.1.100", "172.1.2.3"]
     command = build_smb_command(
-        servers, 2, OperatingSystem.WINDOWS, DROPPER_EXE_PATH, AGENT_EXE_PATH, otp_provider
+        AGENT_ID,
+        servers,
+        2,
+        OperatingSystem.WINDOWS,
+        DROPPER_EXE_PATH,
+        AGENT_EXE_PATH,
+        otp_provider,
     )
 
     for s in servers:
@@ -55,6 +64,7 @@ def test_servers(
 
 def test_dropper_used(otp_provider: IAgentOTPProvider):
     command = build_smb_command(
+        AGENT_ID,
         ["127.0.0.1"],
         2,
         OperatingSystem.WINDOWS,
@@ -68,6 +78,7 @@ def test_dropper_used(otp_provider: IAgentOTPProvider):
 
 def test_monkey_used(otp_provider: IAgentOTPProvider):
     command = build_smb_command(
+        AGENT_ID,
         ["127.0.0.1"],
         2,
         OperatingSystem.WINDOWS,
@@ -81,6 +92,7 @@ def test_monkey_used(otp_provider: IAgentOTPProvider):
 
 def test_otp_used(otp_provider: IAgentOTPProvider):
     command = build_smb_command(
+        AGENT_ID,
         ["127.0.0.1"],
         2,
         OperatingSystem.WINDOWS,

monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py

@@ -7,6 +7,7 @@
 
 from infection_monkey.exploit import powershell
 from infection_monkey.exploit.tools.helpers import AGENT_BINARY_PATH_WIN64
+from infection_monkey.utils.ids import get_agent_id
 
 # Use the path_win32api_get_user_name fixture for all tests in this module
 pytestmark = pytest.mark.usefixtures("patch_win32api_get_user_name")
@@ -18,6 +19,7 @@
 
 bogus_servers = ["1.1.1.1:5000", "2.2.2.2:5007"]
 VICTIM_IP = IPv4Address("10.10.10.1")
+AGENT_ID = get_agent_id()
 
 
 mock_agent_binary_repository = MagicMock()
@@ -53,6 +55,7 @@ def powershell_arguments(host_with_ip_address):
         },
     }
     arguments = {
+        "agent_id": AGENT_ID,
         "host": host_with_ip_address,
         "servers": bogus_servers,
         "options": options,
@@ -186,7 +189,7 @@ def test_build_monkey_execution_command():
     depth = 2
     executable_path = "/tmp/test-monkey"
 
-    cmd = powershell.build_monkey_execution_command(bogus_servers, depth, executable_path)
+    cmd = powershell.build_monkey_execution_command(AGENT_ID, bogus_servers, depth, executable_path)
 
     assert f"-d {depth}" in cmd
     assert executable_path in cmd

monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py

@@ -6,6 +6,7 @@
 from common.agent_events import ExploitationEvent, PasswordRestorationEvent
 from common.event_queue import IAgentEventQueue
 from infection_monkey.i_puppet import TargetHost
+from infection_monkey.utils.ids import get_agent_id
 
 NETBIOS_NAME = "NetBIOS Name"
 
@@ -32,6 +33,7 @@ def mock_report_login_attempt(**kwargs):
     monkeypatch.setattr(obj, "report_login_attempt", mock_report_login_attempt)
     monkeypatch.setattr(obj, "host", TargetHost(ip=IPv4Address("1.1.1.1")))
     monkeypatch.setattr(obj, "agent_event_queue", mock_agent_event_queue)
+    monkeypatch.setattr(obj, "agent_id", get_agent_id())
     return obj
 
 

monkey/tests/unit_tests/infection_monkey/utils/test_commands.py

@@ -9,6 +9,11 @@
 from infection_monkey.utils.ids import get_agent_id
 
 
+@pytest.fixture
+def agent_id():
+    return get_agent_id()
+
+
 def test_build_monkey_commandline_explicitly_arguments():
     expected = [
         "-p",
@@ -73,19 +78,22 @@ def test_get_monkey_commandline_linux():
     assert expected == actual
 
 
-def test_build_monkey_commandline():
+def test_build_monkey_commandline(agent_id):
     servers = ["10.10.10.10:5000", "11.11.11.11:5007"]
 
-    expected = f" -p {get_agent_id()} -s 10.10.10.10:5000,11.11.11.11:5007 -d 0 -l /home/bla"
-    actual = build_monkey_commandline(servers=servers, depth=0, location="/home/bla")
+    expected = f" -p {agent_id} -s 10.10.10.10:5000,11.11.11.11:5007 -d 0 -l /home/bla"
+    actual = build_monkey_commandline(
+        agent_id=agent_id, servers=servers, depth=0, location="/home/bla"
+    )
 
     assert expected == actual
 
 
 @pytest.mark.parametrize("servers", [None, []])
-def test_build_monkey_commandline_empty_servers(servers):
-
-    expected = f" -p {get_agent_id()} -d 0 -l /home/bla"
-    actual = build_monkey_commandline(servers, depth=0, location="/home/bla")
+def test_build_monkey_commandline_empty_servers(agent_id, servers):
+    expected = f" -p {agent_id} -d 0 -l /home/bla"
+    actual = build_monkey_commandline(
+        agent_id=agent_id, servers=servers, depth=0, location="/home/bla"
+    )
 
     assert expected == actual