SNMP options

monkey/agent_plugins/exploiters/snmp/config-schema.json

@@ -0,0 +1,25 @@
+{
+    "properties": {
+        "agent_binary_download_timeout": {
+            "title": "Agent binary download timeout",
+            "description": "The maximum time (in seconds) to wait for a successfully exploited SNMP client to download the Agent binary.",
+            "type": "number",
+            "minimum": 0.0,
+            "default": 60.0
+        },
+        "snmp_request_timeout": {
+            "title": "SNMP request timeout",
+            "description": "The maximum time (in seconds) to wait for a response from an SNMP client.",
+            "type": "number",
+            "minimum": 0.0,
+            "default": 0.5
+        },
+        "snmp_retries": {
+            "title": "SNMP retries",
+            "description": "The number of times to retry an SNMP request before giving up.",
+            "type": "integer",
+            "minimum": 0,
+            "default": 3
+        }
+    }
+}

monkey/agent_plugins/exploiters/snmp/src/plugin.py

@@ -19,7 +19,8 @@
 from .community_string_generator import generate_community_strings
 from .snmp_client import SNMPClient
 from .snmp_exploit_client import SNMPExploitClient
-from .snmp_exploiter import SNMPExploiter, StubSNMPOptions
+from .snmp_exploiter import SNMPExploiter
+from .snmp_options import SNMPOptions
 
 logger = logging.getLogger(__name__)
 
@@ -50,25 +51,13 @@ def __init__(
         otp_provider: IAgentOTPProvider,
         **kwargs,
     ):
-        self._snmp_client = SNMPClient()
-        exploit_client = SNMPExploitClient(
-            agent_id, agent_event_publisher, plugin_name, self._snmp_client
-        )
-        agent_binary_server_factory = partial(
-            start_agent_binary_server,
-            agent_binary_repository=agent_binary_repository,
-            tcp_port_selector=tcp_port_selector,
-        )
-        get_community_strings = partial(
-            generate_community_strings, propagation_credentials_repository
-        )
-        self._snmp_exploiter = SNMPExploiter(
-            agent_id,
-            exploit_client,
-            agent_binary_server_factory,
-            get_community_strings,
-            otp_provider,
-        )
+        self._plugin_name = plugin_name
+        self._agent_id = agent_id
+        self._agent_event_publisher = agent_event_publisher
+        self._agent_binary_repository = agent_binary_repository
+        self._propagation_credentials_repository = propagation_credentials_repository
+        self._tcp_port_selector = tcp_port_selector
+        self._otp_provider = otp_provider
 
     def run(
         self,
@@ -85,15 +74,17 @@ def run(
 
         try:
             logger.debug(f"Parsing options: {pformat(options)}")
-            snmp_options = StubSNMPOptions(**options)
+            snmp_options = SNMPOptions(**options)
         except Exception as err:
             msg = f"Failed to parse SNMP options: {err}"
             logger.exception(msg)
             return ExploiterResultData(
                 exploitation_success=False, propagation_success=False, error_message=msg
             )
 
-        attempt_exploit, msg = should_attempt_exploit(host, self._snmp_client)
+        snmp_client = SNMPClient(snmp_options.snmp_request_timeout, snmp_options.snmp_retries)
+
+        attempt_exploit, msg = should_attempt_exploit(host, snmp_client)
         if not attempt_exploit:
             logger.debug(f"Skipping brute force of host {host.ip}: {msg}")
             return ExploiterResultData(
@@ -102,12 +93,33 @@ def run(
 
         try:
             logger.debug(f"Running SNMP exploiter on host {host.ip}")
-            return self._snmp_exploiter.exploit_host(
-                host, servers, current_depth, snmp_options, interrupt
+            community_strings = generate_community_strings(
+                self._propagation_credentials_repository.get_credentials()
+            )
+
+            snmp_exploiter = self._create_snmp_exploiter(snmp_client)
+            return snmp_exploiter.exploit_host(
+                host, servers, current_depth, snmp_options, community_strings, interrupt
             )
         except Exception as err:
             msg = f"An unexpected exception occurred while attempting to exploit host: {err}"
             logger.exception(msg)
             return ExploiterResultData(
                 exploitation_success=False, propagation_success=False, error_message=msg
             )
+
+    def _create_snmp_exploiter(self, snmp_client: SNMPClient) -> SNMPExploiter:
+        exploit_client = SNMPExploitClient(
+            self._agent_id, self._agent_event_publisher, self._plugin_name, snmp_client
+        )
+        agent_binary_server_factory = partial(
+            start_agent_binary_server,
+            agent_binary_repository=self._agent_binary_repository,
+            tcp_port_selector=self._tcp_port_selector,
+        )
+        return SNMPExploiter(
+            self._agent_id,
+            exploit_client,
+            agent_binary_server_factory,
+            self._otp_provider,
+        )

monkey/agent_plugins/exploiters/snmp/src/snmp_client.py

@@ -32,11 +32,17 @@ class SNMPResult:
 
 
 class SNMPClient:
-    def __init__(self):
+    def __init__(self, snmp_request_timeout: float, snmp_retries: int):
+        """
+        :param snmp_request_timeout: The timeout for SNMP requests, in seconds
+        :param snmp_retries: The number of retries for SNMP requests
+        """
         self._engine = SnmpEngine()
 
         mib_builder = self._engine.getMibBuilder()
         mib_builder.addMibSources(builder.DirMibSource(dirname(__file__)))
+        self._snmp_request_timeout = snmp_request_timeout
+        self._snmp_retries = snmp_retries
 
     def get_system_name(self, target_ip: IPv4Address, community_string: str) -> str:
         """
@@ -51,7 +57,11 @@ def get_system_name(self, target_ip: IPv4Address, community_string: str) -> str:
         cmd = getCmd(
             self._engine,
             CommunityData(community_string),
-            UdpTransportTarget((str(target_ip), 161)),
+            UdpTransportTarget(
+                (str(target_ip), 161),
+                timeout=self._snmp_request_timeout,
+                retries=self._snmp_retries,
+            ),
             ContextData(),
             ObjectType(ObjectIdentity("SNMPv2-MIB", "sysName", 0)),
         )
@@ -80,7 +90,11 @@ def set_command(
         cmd = setCmd(
             self._engine,
             CommunityData(community_string),
-            UdpTransportTarget((str(target_ip), 161)),
+            UdpTransportTarget(
+                (str(target_ip), 161),
+                timeout=self._snmp_request_timeout,
+                retries=self._snmp_retries,
+            ),
             ContextData(),
             ObjectType(
                 ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", f'"{command_name}"'),
@@ -116,7 +130,11 @@ def execute_command(self, target_ip: IPv4Address, command_name: str, community_s
         cmd = getCmd(
             self._engine,
             CommunityData(community_string),
-            UdpTransportTarget((str(target_ip), 161)),
+            UdpTransportTarget(
+                (str(target_ip), 161),
+                timeout=self._snmp_request_timeout,
+                retries=self._snmp_retries,
+            ),
             ContextData(),
             ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendNumEntries", 0)),
             ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", f'"{command_name}"')),
@@ -140,7 +158,11 @@ def clear_command(self, target_ip: IPv4Address, command_name: str, community_str
         cmd = setCmd(
             self._engine,
             CommunityData(community_string),
-            UdpTransportTarget((str(target_ip), 161)),
+            UdpTransportTarget(
+                (str(target_ip), 161),
+                timeout=self._snmp_request_timeout,
+                retries=self._snmp_retries,
+            ),
             ContextData(),
             ObjectType(
                 ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", f'"{command_name}"'),

monkey/agent_plugins/exploiters/snmp/src/snmp_exploit_client.py

@@ -16,6 +16,7 @@
 from infection_monkey.i_puppet import TargetHost
 
 from .snmp_client import SNMPClient
+from .snmp_options import SNMPOptions
 
 COMMAND_NAME_LENGTH = 6
 
@@ -49,7 +50,12 @@ def __init__(
         self._generate_command_name = generate_command_name
 
     def exploit_host(
-        self, host: TargetHost, community_string: str, command: str, agent_binary_downloaded: Event
+        self,
+        host: TargetHost,
+        community_string: str,
+        command: str,
+        agent_binary_downloaded: Event,
+        options: SNMPOptions,
     ) -> Tuple[bool, bool]:
         """
         Exploit the host using SNMP
@@ -72,7 +78,7 @@ def exploit_host(
             exploitation_success = False
 
         propagation_success = self._evaluate_propagation_success(
-            exploitation_success, agent_binary_downloaded, 1
+            exploitation_success, agent_binary_downloaded, options.agent_binary_download_timeout
         )
 
         self._publish_exploitation_event(

monkey/agent_plugins/exploiters/snmp/src/snmp_exploiter.py

@@ -9,16 +9,11 @@
 
 from .snmp_command_builder import build_snmp_command
 from .snmp_exploit_client import SNMPExploitClient
+from .snmp_options import SNMPOptions
 
 logger = getLogger(__name__)
 
 AgentBinaryServerFactory = Callable[[TargetHost], HTTPBytesServer]
-CommunityStringGenerator = Callable[[], Iterable[str]]
-
-
-class StubSNMPOptions:
-    def __init__(self, *args, **kwargs):
-        pass
 
 
 class SNMPExploiter:
@@ -27,21 +22,20 @@ def __init__(
         agent_id: AgentID,
         exploit_client: SNMPExploitClient,
         agent_binary_server_factory: AgentBinaryServerFactory,
-        generate_community_strings: CommunityStringGenerator,
         otp_provider: IAgentOTPProvider,
     ):
         self._agent_id = agent_id
         self._exploit_client = exploit_client
         self._agent_binary_server_factory = agent_binary_server_factory
-        self._generate_community_strings = generate_community_strings
         self._otp_provider = otp_provider
 
     def exploit_host(
         self,
         host: TargetHost,
         servers: Sequence[str],
         current_depth: int,
-        options: StubSNMPOptions,
+        options: SNMPOptions,
+        community_strings: Iterable[str],
         interrupt: Event,
     ) -> ExploiterResultData:
         try:
@@ -70,6 +64,7 @@ def exploit_host(
             return self._brute_force_exploit_host(
                 host,
                 options,
+                community_strings,
                 command,
                 agent_binary_http_server.bytes_downloaded,
                 interrupt,
@@ -86,14 +81,15 @@ def exploit_host(
     def _brute_force_exploit_host(
         self,
         host: TargetHost,
-        options: StubSNMPOptions,
+        options: SNMPOptions,
+        community_strings: Iterable[str],
         command: str,
         agent_binary_downloaded: Event,
         interrupt: Event,
     ) -> ExploiterResultData:
         exploit_result = ExploiterResultData(exploitation_success=False, propagation_success=False)
 
-        for community_string in interruptible_iter(self._generate_community_strings(), interrupt):
+        for community_string in interruptible_iter(community_strings, interrupt):
             (
                 exploit_result.exploitation_success,
                 exploit_result.propagation_success,
@@ -102,6 +98,7 @@ def _brute_force_exploit_host(
                 community_string,
                 command,
                 agent_binary_downloaded,
+                options,
             )
 
             if exploit_result.exploitation_success:

monkey/agent_plugins/exploiters/snmp/src/snmp_options.py

@@ -0,0 +1,22 @@
+from pydantic import Field
+
+from common.base_models import InfectionMonkeyBaseModel
+
+
+class SNMPOptions(InfectionMonkeyBaseModel):
+    agent_binary_download_timeout: float = Field(
+        gt=0.0,
+        default=60.0,
+        description="The maximum time (in seconds) to wait for a successfully exploited"
+        " SNMP client to download the agent binary.",
+    )
+    snmp_request_timeout: float = Field(
+        gt=0.0,
+        default=0.5,
+        description="The maximum time (in seconds) to wait for a response from an SNMP client.",
+    )
+    snmp_retries: int = Field(
+        ge=0,
+        default=3,
+        description="The number of times to retry an SNMP request before giving up.",
+    )

monkey/tests/unit_tests/agent_plugins/exploiters/snmp/test_snmp_client.py

@@ -43,7 +43,7 @@ def mock_snmp_set() -> MagicMock:
 def snmp_client(monkeypatch, mock_snmp_get, mock_snmp_set) -> Callable[[], MagicMock]:
     monkeypatch.setattr("agent_plugins.exploiters.snmp.src.snmp_client.getCmd", mock_snmp_get)
     monkeypatch.setattr("agent_plugins.exploiters.snmp.src.snmp_client.setCmd", mock_snmp_set)
-    return SNMPClient()
+    return SNMPClient(snmp_request_timeout=0.5, snmp_retries=3)
 
 
 @dataclass

monkey/tests/unit_tests/agent_plugins/exploiters/snmp/test_snmp_exploit_client.py

@@ -7,6 +7,7 @@
 import pytest
 from agent_plugins.exploiters.snmp.src.snmp_client import SNMPClient, SNMPRequestTimedOut
 from agent_plugins.exploiters.snmp.src.snmp_exploit_client import SNMPExploitClient
+from agent_plugins.exploiters.snmp.src.snmp_options import SNMPOptions
 
 from common import OperatingSystem
 from common.agent_events import ExploitationEvent, PropagationEvent
@@ -77,6 +78,7 @@ def _inner() -> Tuple[bool, bool]:
             COMMUNITY_STRING,
             COMMAND,
             agent_binary_downloaded,
+            SNMPOptions(agent_binary_download_timeout=0.001),
         )
 
     return _inner