wmi: tell the user when a step hits --exec-timeout or the cleanup fails, and drop a dead stub

wmi -x/-X: run the command, then poll Win32_Process by PID (and its
children) until it exits, bounded by --exec-timeout, instead of fixed
sleeps. certutil and reg add run as one chained process.

nxc/protocols/wmi.py

@@ -722,8 +722,6 @@ class wmi(connection):
 
     @requires_admin
     def execute(self, command=None, get_output=False, use_powershell=False):
-        output = ""
-
         # Execution via -x
         if not command and self.args.execute:
             command = self.args.execute
@@ -756,17 +754,18 @@ class wmi(connection):
                 self.args.codec,
                 self.get_namespace("//./root/subscription")
             )
-        output = exec_method.execute(command, get_output, use_powershell=use_powershell)
+        try:
+            output = exec_method.execute(command, get_output, use_powershell=use_powershell)
+        except Exception as e:
+            self.logger.fail(f"Execute command failed with {self.args.exec_method}: {e}")
+            return ""
 
-        if self.args.execute and get_output:
-            self.logger.success(f'Executed command: "{command}" via {self.args.exec_method}')
-            buf = StringIO(output).readlines()
-            for line in buf:
+        if get_output and (self.args.execute_psh if use_powershell else self.args.execute):
+            self.logger.success(f'Executed {"PowerShell " if use_powershell else ""}command: "{command}" via {self.args.exec_method}')
+            for line in StringIO(output).readlines():
                 if line.strip():
                     self.logger.highlight(line.strip())
-            return output
-        else:
-            return output
+        return output
 
     def execute_psh(self, command=None, get_output=False):
         # Execution via -X
@@ -774,18 +773,7 @@ class wmi(connection):
             command = self.args.execute_psh
             if not self.args.no_output:
                 get_output = True
-
-        output = self.execute(command, get_output, use_powershell=True)
-
-        if self.args.execute_psh and get_output:
-            self.logger.success(f'Executed PowerShell command: "{command}" via {self.args.exec_method}')
-            buf = StringIO(output).readlines()
-            for line in buf:
-                if line.strip():
-                    self.logger.highlight(line.strip())
-            return output
-        else:
-            return output
+        return self.execute(command, get_output, use_powershell=True)
 
     def get_namespace(self, namespace: str) -> IWbemServices:
         """Load WMI namespaces and place them in cache. If a namespace is already loaded in cache, return the namespace in cache"""

nxc/protocols/wmi/wmiexec.py

@@ -57,10 +57,40 @@ class WMIEXEC:
 
     def execute_remote(self, command):
         self.logger.info(f"Executing command: {command}")
+        process = self.__win32Process.Create(command, self.__pwd, None)
+        if process.ReturnValue != 0:
+            raise Exception(f"Win32_Process.Create returned {process.ReturnValue}")
+        return process.ProcessId
+
+    def wait_for_process(self, pid, timeout):
+        """Poll Win32_Process until the process and its children have exited or the timeout (in seconds) is reached"""
+        end = time.monotonic() + timeout
+        while time.monotonic() < end:
+            time.sleep(0.1)
+            try:
+                self.__iWbemServices.GetObject(f'Win32_Process.Handle="{pid}"')
+            except Exception as e:
+                if "WBEM_E_NOT_FOUND" not in str(e):
+                    self.logger.debug(f"Could not query process {pid}: {e}")
+                    continue
+                # Children (e.g. "start /b") keep the ParentProcessId after the parent exits and may still hold the output file open
+                if not self.has_children(pid):
+                    return
+        self.logger.fail(f"WMIEXEC: Process {pid} still running after {timeout}s, output may be incomplete. Please try a higher '--exec-timeout'")
+
+    def has_children(self, pid):
         try:
-            self.__win32Process.Create(command, self.__pwd, None)
+            enum = self.__iWbemServices.ExecQuery(f"SELECT ProcessId FROM Win32_Process WHERE ParentProcessId = {pid}")
         except Exception as e:
-            self.logger.error(str(e))
+            self.logger.debug(f"Could not query children of process {pid}: {e}")
+            return False
+        try:
+            enum.Next(0xFFFFFFFF, 1)
+            return True
+        except Exception:
+            return False
+        finally:
+            enum.RemRelease()
 
     def execute_WithOutput(self, command):
         result_output = f"C:\\windows\\temp\\{uuid.uuid4()!s}.txt"
@@ -69,17 +99,11 @@ class WMIEXEC:
         self.__registry_Path = f"Software\\Classes\\{gen_random_string(8)}"
 
         # 1. Run the command and write output to file
-        self.execute_remote(f'{self.__shell} ({command}) 1> "{result_output}" 2>&1')
-        self.logger.info(f"Waiting {self.__exec_timeout}s for command to complete.")
-        time.sleep(self.__exec_timeout)
+        self.logger.info(f"Waiting up to {self.__exec_timeout}s for command to complete.")
+        self.wait_for_process(self.execute_remote(f'{self.__shell} ({command}) 1> "{result_output}" 2>&1'), self.__exec_timeout)
 
-        # 2. Base64 encode the file
-        self.execute_remote(f"{self.__shell} certutil -encodehex -f {result_output} {result_output_b64} 0x40000001")
-        time.sleep(0.5)
-
-        # 3. Store content in registry
-        self.execute_remote(f'{self.__shell} for /F "usebackq" %G in ("{result_output_b64}") do reg add HKLM\\{self.__registry_Path} /v {keyName} /t REG_SZ /d "%G" /f')
-        time.sleep(0.1)
+        # 2. Base64 encode the file and store content in registry
+        self.wait_for_process(self.execute_remote(f'{self.__shell} certutil -encodehex -f {result_output} {result_output_b64} 0x40000001 & for /F "usebackq" %G in ("{result_output_b64}") do reg add HKLM\\{self.__registry_Path} /v {keyName} /t REG_SZ /d "%G" /f'), self.__exec_timeout)
 
         self.queryRegistry(keyName)
         self.clean_up(result_output, result_output_b64)
@@ -112,16 +136,14 @@ class WMIEXEC:
         # 1. Run the command and write output to file
         if not command.lower().startswith("powershell"):
             command = f'powershell -Command "& {{{command}}}"'
-        self.execute_remote(f'{command} > "{result_output}" 2>&1')
-        self.logger.info(f"Waiting {self.__exec_timeout}s for command to complete.")
-        time.sleep(self.__exec_timeout)
+        self.logger.info(f"Waiting up to {self.__exec_timeout}s for command to complete.")
+        self.wait_for_process(self.execute_remote(f'{command} > "{result_output}" 2>&1'), self.__exec_timeout)
 
         # 2. Base64 encode the file using PowerShell
-        self.execute_remote(f'powershell -Command "[Convert]::ToBase64String([IO.File]::ReadAllBytes(\'{result_output}\')) | Out-File -Encoding ASCII \'{result_output_b64}\'"')
-        time.sleep(0.5)
+        self.wait_for_process(self.execute_remote(f'powershell -Command "[Convert]::ToBase64String([IO.File]::ReadAllBytes(\'{result_output}\')) | Out-File -Encoding ASCII \'{result_output_b64}\'"'), self.__exec_timeout)
 
         # 3. Use PowerShell to split base64 content into 16KB chunks and store in registry
-        self.execute_remote(
+        self.wait_for_process(self.execute_remote(
             f'powershell -Command "$b64 = Get-Content -Raw \'{result_output_b64}\'; '
             f'$chunksize = 16000; '
             f'$count = [math]::Ceiling($b64.Length / $chunksize); '
@@ -130,8 +152,7 @@ class WMIEXEC:
             f'  $name = \\"{keyName}_chunk_$i\\"; '
             f'  reg add \\"HKLM\\{self.__registry_Path}\\" /v $name /t REG_SZ /d $chunk /f }}; '
             f'reg add \\"HKLM\\{self.__registry_Path}\\" /v \\"{keyName}\\" /t REG_DWORD /d $count /f"'
-        )
-        time.sleep(0.1)
+        ), self.__exec_timeout)
 
         self.queryRegistry_psh(keyName)
         self.clean_up(result_output, result_output_b64)
@@ -165,13 +186,12 @@ class WMIEXEC:
             self.logger.fail("WMIEXEC: Could not retrieve output file! Either command timed out or AV killed the process. Please try increasing the timeout: '--exec-timeout 10'")
 
     def clean_up(self, result_output, result_output_b64):
-        """Deletes the output file, the base64 output file, and the registry path where the base64 content was stored."""
-        self.execute_remote(f'{self.__shell} del /q /f "{result_output}" "{result_output_b64}"')
-
+        """Deletes the registry path where the base64 content was stored, the output file and the base64 output file."""
         try:
             self.logger.debug(f"Removing temporary registry path: HKLM\\{self.__registry_Path}")
             descriptor, _ = self.__iWbemServices.GetObject("StdRegProv")
             descriptor = descriptor.SpawnInstance()
             descriptor.DeleteKey(0x80000002, self.__registry_Path)
+            self.execute_remote(f'{self.__shell} del /q /f "{result_output}" "{result_output_b64}"')
         except Exception as e:
-            self.logger.fail(f"Target: {self.__target} removing temporary registry path error: {e!s}")
+            self.logger.fail(f"WMIEXEC: Cleanup failed, traces may be left on the target ({result_output}, {result_output_b64}, HKLM\\{self.__registry_Path}): {e}")