Agent: Shorten reservation ID

Per Issue #4187, there are cases when download strings must be shorter.
In order to achieve this, random strings from a set of 62 characters are
generated. Using 5 characters gives something like 916M possible values,
which is more than enough for any single agent.

Issue #4187
PR #4189

Author: mssalvatore

monkey/infection_monkey/Pipfile

@@ -26,7 +26,7 @@ email-validator = "*"
 monkey-types = "*"
 monkeyevents = "*"
 monkeytoolbox = "*"
-monkey-agentpluginapi = ">=0.7.0"
+monkey-agentpluginapi = "*"
 
 [dev-packages]
 mypy = "*"

monkey/infection_monkey/Pipfile.lock

@@ -3,7 +3,6 @@
 from http.server import HTTPServer
 from ipaddress import IPv4Address
 from typing import Callable, Optional, Type
-from uuid import uuid4
 
 from agentpluginapi import (
     AgentBinaryDownloadReservation,
@@ -12,7 +11,11 @@
     LocalMachineInfo,
     ReservationID,
 )
-from monkeytoolbox import create_daemon_thread, insecure_generate_random_string
+from monkeytoolbox import (
+    create_daemon_thread,
+    insecure_generate_random_string,
+    secure_generate_random_string,
+)
 from monkeytypes import Event, Lock, NetworkPort, OperatingSystem
 
 from .http_agent_binary_request_handler import AgentBinaryHTTPRequestHandler
@@ -79,7 +82,7 @@ def register(
             if not self.server_is_running():
                 self._start_server()
 
-            reservation_id = uuid4()
+            reservation_id = secure_generate_random_string(n=5)
             url = self._build_request_url(reservation_id, operating_system, requestor_ip)
             reservation = AgentBinaryDownloadReservation(
                 reservation_id,

monkey/infection_monkey/exploit/http_agent_binary_server.py

@@ -2,7 +2,7 @@
 from http import HTTPStatus
 from http.server import HTTPServer
 from io import BytesIO
-from typing import Type
+from typing import Final, Type
 from unittest.mock import MagicMock
 
 import pytest
@@ -20,16 +20,15 @@
     get_http_handler,
 )
 
-AGENT_BINARY = b"agent_binary"
-DROPPER_BINARY = b"dropper_agent_binary"
-IP = "127.0.0.1"
-UUID_1 = ReservationID("00000000-0000-0000-0000-000000000001")
-UUID_2 = ReservationID("00000000-0000-0000-0000-000000000002")
-UUID_3 = ReservationID("00000000-0000-0000-0000-000000000003")
+AGENT_BINARY: Final = b"agent_binary"
+DROPPER_BINARY: Final = b"dropper_agent_binary"
+IP: Final = "127.0.0.1"
+RESERVATION_ID_1: Final = ReservationID("abcABC1")
+RESERVATION_ID_2: Final = ReservationID("abcABC2")
+RESERVATION_ID_3: Final = ReservationID("abcABC2")
 
-
-DEFAULT_AGENT_TEMPLATE = b"%(agent_binary)s"
-DROPPER_AGENT_TEMPLATE = b"dropper_%(agent_binary)s"
+DEFAULT_AGENT_TEMPLATE: Final = b"%(agent_binary)s"
+DROPPER_AGENT_TEMPLATE: Final = b"dropper_%(agent_binary)s"
 
 
 @pytest.fixture
@@ -40,54 +39,54 @@ def port(tcp_port_selector) -> int:
 @pytest.fixture
 def binary_request_1(port) -> AgentBinaryDownloadReservation:
     return AgentBinaryDownloadReservation(
-        UUID_1,
+        RESERVATION_ID_1,
         OperatingSystem.LINUX,
         DEFAULT_AGENT_TEMPLATE,
-        f"http://{IP}:{port}/{UUID_1}",
+        f"http://{IP}:{port}/{RESERVATION_ID_1}",
         threading.Event(),
     )
 
 
 @pytest.fixture
 def binary_request_2(port) -> AgentBinaryDownloadReservation:
     return AgentBinaryDownloadReservation(
-        UUID_2,
+        RESERVATION_ID_2,
         OperatingSystem.WINDOWS,
         DEFAULT_AGENT_TEMPLATE,
-        f"http://{IP}:{port}/{UUID_2}",
+        f"http://{IP}:{port}/{RESERVATION_ID_2}",
         threading.Event(),
     )
 
 
 @pytest.fixture
 def binary_request_3(port) -> AgentBinaryDownloadReservation:
     return AgentBinaryDownloadReservation(
-        UUID_2,
+        RESERVATION_ID_2,
         OperatingSystem.WINDOWS,
         None,
-        f"http://{IP}:{port}/{UUID_2}",
+        f"http://{IP}:{port}/{RESERVATION_ID_2}",
         threading.Event(),
     )
 
 
 @pytest.fixture
 def dropper_request_1(port) -> AgentBinaryDownloadReservation:
     return AgentBinaryDownloadReservation(
-        UUID_1,
+        RESERVATION_ID_1,
         OperatingSystem.LINUX,
         DROPPER_AGENT_TEMPLATE,
-        f"http://{IP}:{port}/{UUID_1}",
+        f"http://{IP}:{port}/{RESERVATION_ID_1}",
         threading.Event(),
     )
 
 
 @pytest.fixture
 def dropper_request_2(port) -> AgentBinaryDownloadReservation:
     return AgentBinaryDownloadReservation(
-        UUID_2,
+        RESERVATION_ID_2,
         OperatingSystem.WINDOWS,
         DROPPER_AGENT_TEMPLATE,
-        f"http://{IP}:{port}/{UUID_2}",
+        f"http://{IP}:{port}/{RESERVATION_ID_2}",
         threading.Event(),
     )
 
@@ -302,3 +301,4 @@ def test_agent_binary_request__is_transformed(
 
     assert response.status_code == HTTPStatus.OK
     assert response.content == DROPPER_BINARY
+    assert response.content == DROPPER_BINARY

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

@@ -7,7 +7,7 @@
 from multiprocessing.managers import SyncManager
 from pathlib import Path
 from queue import Queue
-from typing import List, Tuple, Type
+from typing import Final, List, Tuple, Type
 from unittest.mock import MagicMock
 
 import pytest
@@ -19,8 +19,8 @@
 from infection_monkey.exploit.http_agent_binary_server import HTTPAgentBinaryServer
 from infection_monkey.network import TCPPortSelector
 
-REQUESTOR_IP = IPv4Address("1.1.1.1")
-UUID_1 = ReservationID("00000000-0000-0000-0000-000000000001")
+REQUESTOR_IP: Final = IPv4Address("1.1.1.1")
+RESERVATION_ID_1: Final = ReservationID("abcdABCD1")
 
 
 def use_agent_binary(agent_binary: bytes) -> bytes:
@@ -188,7 +188,7 @@ def test_deregister__raises_error_on_invalid_reservation_id(
     mock_http_handler = mock_agent_binary_http_handler
     mock_http_handler.clear_reservation_mock.side_effect = KeyError  # type: ignore[attr-defined]
     with pytest.raises(KeyError):
-        http_agent_binary_server.deregister(UUID_1)
+        http_agent_binary_server.deregister(RESERVATION_ID_1)
 
 
 @pytest.mark.xdist_group(name="tcp_port_selector")