From 130c62a5c2448339723d61c80362e8cf35adc163 Mon Sep 17 00:00:00 2001 From: vakaris_zilius Date: Wed, 9 Mar 2022 14:17:36 +0000 Subject: [PATCH] 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 | 13 +++++++++++-- monkey/infection_monkey/exploit/wmiexec.py | 1 + 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/exploit/tools/wmi_tools.py b/monkey/infection_monkey/exploit/tools/wmi_tools.py index 99fda3141..30ae59107 100644 --- a/monkey/infection_monkey/exploit/tools/wmi_tools.py +++ b/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,8 +93,7 @@ class WmiTools(object): def dcom_wrap(func): def _wrapper(*args, **kwarg): try: - with lock: - return func(*args, **kwarg) + return func(*args, **kwarg) finally: WmiTools.dcom_cleanup() diff --git a/monkey/infection_monkey/exploit/wmiexec.py b/monkey/infection_monkey/exploit/wmiexec.py index de7ab58c7..a81877df1 100644 --- a/monkey/infection_monkey/exploit/wmiexec.py +++ b/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: