Compare commits

...
3 changed files with 66 additions and 11 deletions
+1
View File
@@ -50,6 +50,7 @@ All notable changes to the **Prowler SDK** are documented in this file.
- Check `check_name` has no `resource_name` error for GCP provider [(#9169)](https://github.com/prowler-cloud/prowler/pull/9169) - Check `check_name` has no `resource_name` error for GCP provider [(#9169)](https://github.com/prowler-cloud/prowler/pull/9169)
- Depth Truncation and parsing error in PowerShell queries [(#9181)](https://github.com/prowler-cloud/prowler/pull/9181) - Depth Truncation and parsing error in PowerShell queries [(#9181)](https://github.com/prowler-cloud/prowler/pull/9181)
- Fix M365 Teams `--sp-env-auth` connection error and enhanced timeout logging [(#9191)](https://github.com/prowler-cloud/prowler/pull/9191) - Fix M365 Teams `--sp-env-auth` connection error and enhanced timeout logging [(#9191)](https://github.com/prowler-cloud/prowler/pull/9191)
- Fix M365 Teams connection error and enhanced timeout logging [(#9197)](https://github.com/prowler-cloud/prowler/pull/9197)
--- ---
+28 -5
View File
@@ -3,7 +3,7 @@ import queue
import re import re
import subprocess import subprocess
import threading import threading
from typing import Union from typing import Optional, Union
from prowler.lib.logger import logger from prowler.lib.logger import logger
@@ -55,6 +55,7 @@ class PowerShellSession:
text=True, text=True,
bufsize=1, bufsize=1,
) )
self._last_command_timed_out = False
def sanitize(self, credential: str) -> str: 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-Output '{self.END}'\n")
self.process.stdin.write(f"Write-Error '{self.END}'\n") self.process.stdin.write(f"Write-Error '{self.END}'\n")
return ( return (
self.json_parse_output(self.read_output(timeout=timeout)) self.json_parse_output(self.read_output(timeout=timeout, command=command))
if json_parse 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. Read output from a process with timeout functionality.
@@ -154,6 +157,7 @@ class PowerShellSession:
result_queue = queue.Queue() result_queue = queue.Queue()
error_lines = [] error_lines = []
error_queue = queue.Queue() error_queue = queue.Queue()
self._last_command_timed_out = False
def reader_thread(): def reader_thread():
try: try:
@@ -190,13 +194,32 @@ class PowerShellSession:
result = result_queue.get(timeout=timeout) or default result = result_queue.get(timeout=timeout) or default
error_result = error_queue.get(timeout=1) error_result = error_queue.get(timeout=1)
except queue.Empty: 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: if error_result:
logger.error(f"PowerShell error output: {error_result}") logger.error(f"PowerShell error output: {error_result}")
return 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: def json_parse_output(self, output: str) -> dict:
""" """
Parse command execution output to JSON format. Parse command execution output to JSON format.
+37 -6
View File
@@ -1,3 +1,4 @@
import time
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
from prowler.providers.m365.lib.powershell.m365_powershell import PowerShellSession 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("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-Output '{session.END}'\n")
mock_process.stdin.write.assert_any_call(f"Write-Error '{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 # Test 2: JSON parsing enabled
mock_process.stdout.readline.side_effect = [ mock_process.stdout.readline.side_effect = [
@@ -94,10 +96,22 @@ class TestPowerShellSession:
mock_json_parse.assert_called_once_with('{"key": "value"}') mock_json_parse.assert_called_once_with('{"key": "value"}')
# Test 3: Timeout handling # 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" mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n"
result = session.execute("Get-Command", timeout=0.1) with patch("prowler.lib.logger.logger.error") as mock_error:
assert result == "" 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 # Test 4: Error handling
mock_process.stdout.readline.side_effect = ["\n", f"{session.END}\n"] 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): with patch.object(session, "remove_ansi", side_effect=lambda x: x):
result = session.read_output() result = session.read_output()
assert result == "Hello World" assert result == "Hello World"
assert session.last_command_timed_out is False
# Test 2: Error in stderr # Test 2: Error in stderr
mock_process.stdout.readline.side_effect = ["\n", f"{session.END}\n"] mock_process.stdout.readline.side_effect = ["\n", f"{session.END}\n"]
@@ -148,18 +163,34 @@ class TestPowerShellSession:
mock_error.assert_called_once_with( mock_error.assert_called_once_with(
"PowerShell error output: Write-Error: This is an error" "PowerShell error output: Write-Error: This is an error"
) )
assert session.last_command_timed_out is False
# Test 3: Timeout in stdout # 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" mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n"
result = session.read_output(timeout=0.1, default="timeout") with patch("prowler.lib.logger.logger.error") as mock_error:
assert result == "timeout" 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 # Test 4: Empty output
mock_process.stdout.readline.side_effect = [f"{session.END}\n"] mock_process.stdout.readline.side_effect = [f"{session.END}\n"]
mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n" mock_process.stderr.readline.return_value = f"Write-Error: {session.END}\n"
result = session.read_output() result = session.read_output()
assert result == "" assert result == ""
assert session.last_command_timed_out is False
session.close() session.close()