Agent: add a wrapper for wmi_tools users

Add a dedicated wrapper to make sure that wmi_tools users don't run into race conditions

monkey/infection_monkey/exploit/tools/wmi_tools.py

@@ -10,6 +10,7 @@ logger = logging.getLogger(__name__)
 
 
 # Due to the limitations of impacket library we should only run one WmiConnection at a time
+# Use impacket_user decorator to ensure that no race conditions are happening
 # See comments in https://github.com/guardicore/monkey/pull/1766
 lock = threading.Lock()
 
@@ -23,6 +24,15 @@ class AccessDeniedException(Exception):
 
 
 class WmiTools(object):
+
+    @staticmethod
+    def impacket_user(func):
+        def _wrapper(*args, **kwarg):
+            with lock:
+                return func(*args, **kwarg)
+
+        return _wrapper
+
     class WmiConnection(object):
         def __init__(self):
             self._dcom = None
@@ -83,7 +93,6 @@ class WmiTools(object):
     def dcom_wrap(func):
         def _wrapper(*args, **kwarg):
             try:
-                with lock:
                 return func(*args, **kwarg)
             finally:
                 WmiTools.dcom_cleanup()

monkey/infection_monkey/exploit/wmiexec.py

@@ -24,6 +24,7 @@ class WmiExploiter(HostExploiter):
     EXPLOIT_TYPE = ExploitType.BRUTE_FORCE
     _EXPLOITED_SERVICE = "WMI (Windows Management Instrumentation)"
 
+    @WmiTools.impacket_user
     @WmiTools.dcom_wrap
     def _exploit_host(self) -> ExploiterResultData: