Merge pull request #1839 from guardicore/1650-signed-scripts-timeout

Add timeout to signed script PBA

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

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):

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))