From 4cc57f123695fe7d3913e73d281a3406915a5a08 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 31 Mar 2022 17:29:16 +0530 Subject: [PATCH] Agent: Add timeouts to signed script PBA --- .../post_breach/actions/use_signed_scripts.py | 4 +++- .../signed_script_proxy/signed_script_proxy.py | 9 +++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py b/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py index 470e07bb1..a9224a977 100644 --- a/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py +++ b/monkey/infection_monkey/post_breach/actions/use_signed_scripts.py @@ -3,6 +3,7 @@ import subprocess from typing import Dict from common.common_consts.post_breach_consts import POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC +from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT, SHORT_REQUEST_TIMEOUT from infection_monkey.post_breach.pba import PBA from infection_monkey.post_breach.signed_script_proxy.signed_script_proxy import ( cleanup_changes, @@ -21,6 +22,7 @@ class SignedScriptProxyExecution(PBA): telemetry_messenger, POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC, windows_cmd=" ".join(windows_cmds), + timeout=MEDIUM_REQUEST_TIMEOUT, ) def run(self, options: Dict): @@ -28,7 +30,7 @@ class SignedScriptProxyExecution(PBA): try: if is_windows_os(): original_comspec = subprocess.check_output( # noqa: DUO116 - "if defined COMSPEC echo %COMSPEC%", shell=True + "if defined COMSPEC echo %COMSPEC%", shell=True, timeout=SHORT_REQUEST_TIMEOUT ).decode() super().run(options) return self.pba_data diff --git a/monkey/infection_monkey/post_breach/signed_script_proxy/signed_script_proxy.py b/monkey/infection_monkey/post_breach/signed_script_proxy/signed_script_proxy.py index 12343d8cf..56bae8b25 100644 --- a/monkey/infection_monkey/post_breach/signed_script_proxy/signed_script_proxy.py +++ b/monkey/infection_monkey/post_breach/signed_script_proxy/signed_script_proxy.py @@ -1,5 +1,6 @@ import subprocess +from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from infection_monkey.post_breach.signed_script_proxy.windows.signed_script_proxy import ( get_windows_commands_to_delete_temp_comspec, get_windows_commands_to_proxy_execution_using_signed_script, @@ -16,6 +17,10 @@ def get_commands_to_proxy_execution_using_signed_script(): def cleanup_changes(original_comspec): if is_windows_os(): subprocess.run( # noqa: DUO116 - get_windows_commands_to_reset_comspec(original_comspec), shell=True + get_windows_commands_to_reset_comspec(original_comspec), + shell=True, + timeout=SHORT_REQUEST_TIMEOUT, + ) + subprocess.run( # noqa: DUO116 + get_windows_commands_to_delete_temp_comspec(), shell=True, timeout=SHORT_REQUEST_TIMEOUT ) - subprocess.run(get_windows_commands_to_delete_temp_comspec(), shell=True) # noqa: DUO116