From c401104a6167ca9b824f4c613b4de57b4b2ddb47 Mon Sep 17 00:00:00 2001 From: Daniel Barranquero Date: Mon, 10 Nov 2025 11:25:29 +0100 Subject: [PATCH] fix(powershell): improve timeout errors --- prowler/lib/powershell/powershell.py | 33 ++++++++++++++++--- tests/lib/powershell/powershell_test.py | 43 +++++++++++++++++++++---- 2 files changed, 65 insertions(+), 11 deletions(-) diff --git a/prowler/lib/powershell/powershell.py b/prowler/lib/powershell/powershell.py index 8142fa8e45..c38a0710a4 100644 --- a/prowler/lib/powershell/powershell.py +++ b/prowler/lib/powershell/powershell.py @@ -3,7 +3,7 @@ import queue import re import subprocess import threading -from typing import Union +from typing import Optional, Union from prowler.lib.logger import logger @@ -55,6 +55,7 @@ class PowerShellSession: text=True, bufsize=1, ) + self._last_command_timed_out = False def sanitize(self, credential: str) -> str: """ @@ -121,12 +122,14 @@ class PowerShellSession: self.process.stdin.write(f"Write-Output '{self.END}'\n") self.process.stdin.write(f"Write-Error '{self.END}'\n") return ( - self.json_parse_output(self.read_output(timeout=timeout)) + self.json_parse_output(self.read_output(timeout=timeout, command=command)) if json_parse - else self.read_output(timeout=timeout) + else self.read_output(timeout=timeout, command=command) ) - def read_output(self, timeout: int = 10, default: str = "") -> str: + def read_output( + self, timeout: int = 10, default: str = "", command: Optional[str] = None + ) -> str: """ Read output from a process with timeout functionality. @@ -154,6 +157,7 @@ class PowerShellSession: result_queue = queue.Queue() error_lines = [] error_queue = queue.Queue() + self._last_command_timed_out = False def reader_thread(): try: @@ -190,13 +194,32 @@ class PowerShellSession: result = result_queue.get(timeout=timeout) or default error_result = error_queue.get(timeout=1) except queue.Empty: - result = default + self._last_command_timed_out = True + command_info = f": {command}" if command else "" + logger.error( + f"PowerShell command timed out after {timeout} seconds{command_info}" + ) + thread.join(timeout=timeout) + error_thread.join(timeout=1) + try: + result = result_queue.get_nowait() or default + except queue.Empty: + result = default + try: + error_result = error_queue.get_nowait() + except queue.Empty: + error_result = None if error_result: logger.error(f"PowerShell error output: {error_result}") return result + @property + def last_command_timed_out(self) -> bool: + """Return whether the previous PowerShell command hit the timeout.""" + return self._last_command_timed_out + def json_parse_output(self, output: str) -> dict: """ Parse command execution output to JSON format. diff --git a/tests/lib/powershell/powershell_test.py b/tests/lib/powershell/powershell_test.py index e96740b1d9..8ce20a7e42 100644 --- a/tests/lib/powershell/powershell_test.py +++ b/tests/lib/powershell/powershell_test.py @@ -1,3 +1,4 @@ +import time from unittest.mock import MagicMock, patch from prowler.providers.m365.lib.powershell.m365_powershell import PowerShellSession @@ -78,6 +79,7 @@ class TestPowerShellSession: mock_process.stdin.write.assert_any_call("Get-Command\n") mock_process.stdin.write.assert_any_call(f"Write-Output '{session.END}'\n") mock_process.stdin.write.assert_any_call(f"Write-Error '{session.END}'\n") + assert session.last_command_timed_out is False # Test 2: JSON parsing enabled mock_process.stdout.readline.side_effect = [ @@ -94,10 +96,22 @@ class TestPowerShellSession: mock_json_parse.assert_called_once_with('{"key": "value"}') # Test 3: Timeout handling - mock_process.stdout.readline.side_effect = ["test output\n"] # No END marker + def slow_readline_execute(): + if not hasattr(slow_readline_execute, "called"): + slow_readline_execute.called = True + return "test output\n" + time.sleep(0.2) + return "" + + mock_process.stdout.readline.side_effect = slow_readline_execute mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n" - result = session.execute("Get-Command", timeout=0.1) - assert result == "" + with patch("prowler.lib.logger.logger.error") as mock_error: + result = session.execute("Get-Command", timeout=0.1) + assert result == "" + assert session.last_command_timed_out is True + mock_error.assert_called_once_with( + "PowerShell command timed out after 0.1 seconds: Get-Command" + ) # Test 4: Error handling mock_process.stdout.readline.side_effect = ["\n", f"{session.END}\n"] @@ -134,6 +148,7 @@ class TestPowerShellSession: with patch.object(session, "remove_ansi", side_effect=lambda x: x): result = session.read_output() assert result == "Hello World" + assert session.last_command_timed_out is False # Test 2: Error in stderr mock_process.stdout.readline.side_effect = ["\n", f"{session.END}\n"] @@ -148,18 +163,34 @@ class TestPowerShellSession: mock_error.assert_called_once_with( "PowerShell error output: Write-Error: This is an error" ) + assert session.last_command_timed_out is False # Test 3: Timeout in stdout - mock_process.stdout.readline.side_effect = ["test output\n"] # No END marker + def slow_readline_output(): + if not hasattr(slow_readline_output, "called"): + slow_readline_output.called = True + return "test output\n" + time.sleep(0.2) + return "" + + mock_process.stdout.readline.side_effect = slow_readline_output mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n" - result = session.read_output(timeout=0.1, default="timeout") - assert result == "timeout" + with patch("prowler.lib.logger.logger.error") as mock_error: + result = session.read_output( + timeout=0.1, default="timeout", command="SlowCmd" + ) + assert result == "timeout" + assert session.last_command_timed_out is True + mock_error.assert_called_once_with( + "PowerShell command timed out after 0.1 seconds: SlowCmd" + ) # Test 4: Empty output mock_process.stdout.readline.side_effect = [f"{session.END}\n"] mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n" result = session.read_output() assert result == "" + assert session.last_command_timed_out is False session.close()