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/pba.py b/monkey/infection_monkey/post_breach/pba.py index ba027972e..0ef8e0ecb 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -2,6 +2,7 @@ import logging import subprocess from typing import Dict, Iterable +from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT from common.utils.attack_utils import ScanStatus from infection_monkey.i_puppet.i_puppet import PostBreachData from infection_monkey.telemetry.attack.t1064_telem import T1064Telem @@ -18,7 +19,12 @@ class PBA: """ def __init__( - self, telemetry_messenger: ITelemetryMessenger, name="unknown", linux_cmd="", windows_cmd="" + self, + telemetry_messenger: ITelemetryMessenger, + name="unknown", + linux_cmd="", + windows_cmd="", + timeout: int = LONG_REQUEST_TIMEOUT, ): """ :param name: Name of post breach action. @@ -29,6 +35,7 @@ class PBA: self.name = name self.pba_data = [] self.telemetry_messenger = telemetry_messenger + self.timeout = timeout def run(self, options: Dict) -> Iterable[PostBreachData]: """ @@ -73,12 +80,13 @@ class PBA: """ try: output = subprocess.check_output( # noqa: DUO116 - self.command, stderr=subprocess.STDOUT, shell=True + self.command, stderr=subprocess.STDOUT, shell=True, timeout=self.timeout ).decode() return output, True - except subprocess.CalledProcessError as e: - # Return error output of the command - return e.output.decode(), False + except subprocess.CalledProcessError as err: + return err.output.decode(), False + except subprocess.TimeoutExpired as err: + return str(err), False @staticmethod def choose_command(linux_cmd, windows_cmd): 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..b172d1ab1 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,7 @@ +import logging 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, @@ -7,6 +9,8 @@ from infection_monkey.post_breach.signed_script_proxy.windows.signed_script_prox ) from infection_monkey.utils.environment import is_windows_os +logger = logging.getLogger(__name__) + def get_commands_to_proxy_execution_using_signed_script(): windows_cmds = get_windows_commands_to_proxy_execution_using_signed_script() @@ -15,7 +19,18 @@ 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 - ) - subprocess.run(get_windows_commands_to_delete_temp_comspec(), shell=True) # noqa: DUO116 + try: + subprocess.run( # noqa: DUO116 + 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, + ) + except subprocess.CalledProcessError as err: + logger.error(err.output.decode()) + except subprocess.TimeoutExpired as err: + logger.error(str(err))