Merge pull request #1841 from guardicore/1650-shell-startup-modification

Agent: Add timeouts in shell startup modification PBA's

monkey/infection_monkey/post_breach/actions/modify_shell_startup_files.py

@@ -2,6 +2,7 @@ import subprocess
 from typing import Dict
 
 from common.common_consts.post_breach_consts import POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION
+from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
 from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.post_breach.pba import PBA
 from infection_monkey.post_breach.shell_startup_files.shell_startup_files_modification import (
@@ -21,7 +22,7 @@ class ModifyShellStartupFiles(PBA):
         super().__init__(telemetry_messenger, name=POST_BREACH_SHELL_STARTUP_FILE_MODIFICATION)
 
     def run(self, options: Dict):
-        results = [pba.run() for pba in self.modify_shell_startup_PBA_list()]
+        results = [pba.run(options) for pba in self.modify_shell_startup_PBA_list()]
         if not results:
             results = [
                 (
@@ -70,14 +71,19 @@ class ModifyShellStartupFiles(PBA):
                     windows_cmd=windows_cmds,
                 )
 
-            def run(self):
+            def run(self, options):
                 if self.command:
                     try:
                         output = subprocess.check_output(  # noqa: DUO116
-                            self.command, stderr=subprocess.STDOUT, shell=True
+                            self.command,
+                            stderr=subprocess.STDOUT,
+                            shell=True,
+                            timeout=LONG_REQUEST_TIMEOUT,
                         ).decode()
 
                         return output, True
-                    except subprocess.CalledProcessError as e:
+                    except subprocess.CalledProcessError as err:
                         # Return error output of the command
-                        return e.output.decode(), False
+                        return str(err), False
+                    except subprocess.TimeoutExpired as err:
+                        return str(err), False

monkey/infection_monkey/post_breach/shell_startup_files/linux/shell_startup_files_modification.py

@@ -1,26 +1,37 @@
+import logging
 import subprocess
 
+from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
 from infection_monkey.utils.environment import is_windows_os
 
+logger = logging.getLogger(__name__)
+
 
 def get_linux_commands_to_modify_shell_startup_files():
     if is_windows_os():
         return "", [], []
 
-    HOME_DIR = "/home/"
+    home_dir = "/home/"
+    command = "cut -d: -f1,3 /etc/passwd | egrep ':[0-9]{4}$' | cut -d: -f1"
 
     # get list of usernames
-    USERS = (
+    try:
-        subprocess.check_output(  # noqa: DUO116
+        users = (
-            "cut -d: -f1,3 /etc/passwd | egrep ':[0-9]{4}$' | cut -d: -f1", shell=True
+            subprocess.check_output(  # noqa: DUO116
+                command,
+                shell=True,
+                timeout=MEDIUM_REQUEST_TIMEOUT,
+            )
+            .decode()
+            .split("\n")[:-1]
         )
-        .decode()
+    except subprocess.TimeoutExpired:
-        .split("\n")[:-1]
+        logger.error(f"Command {command} timed out")
-    )
+        return "", [], []
 
     # get list of paths of different shell startup files with place for username
-    STARTUP_FILES = [
+    startup_files = [
-        file_path.format(HOME_DIR)
+        file_path.format(home_dir)
         for file_path in [
             "{0}{{0}}/.profile",  # bash, dash, ksh, sh
             "{0}{{0}}/.bashrc",  # bash
@@ -42,6 +53,6 @@ def get_linux_commands_to_modify_shell_startup_files():
             "tee -a {0} &&",  # append to file
             "sed -i '$d' {0}",  # remove last line of file (undo changes)
         ],
-        STARTUP_FILES,
+        startup_files,
-        USERS,
+        users,
     )

monkey/infection_monkey/post_breach/shell_startup_files/windows/shell_startup_files_modification.py

@@ -1,36 +1,47 @@
+import logging
 import subprocess
 from pathlib import Path
 
+from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
 from infection_monkey.utils.environment import is_windows_os
 
 MODIFY_POWERSHELL_STARTUP_SCRIPT = Path(__file__).parent / "modify_powershell_startup_file.ps1"
 
+logger = logging.getLogger(__name__)
+
 
 def get_windows_commands_to_modify_shell_startup_files():
     if not is_windows_os():
         return "", []
 
     # get powershell startup file path
-    SHELL_STARTUP_FILE = subprocess.check_output("powershell $Profile").decode().split("\r\n")[0]
+    shell_startup_file = subprocess.check_output("powershell $Profile").decode().split("\r\n")[0]
-    SHELL_STARTUP_FILE_PATH_COMPONENTS = SHELL_STARTUP_FILE.split("\\")
+    shell_startup_file_path_components = shell_startup_file.split("\\")
 
     # get list of usernames
-    USERS = (
+    command = "dir C:\\Users /b"
-        subprocess.check_output("dir C:\\Users /b", shell=True)  # noqa: DUO116
+    try:
-        .decode()
+        users = (
-        .split("\r\n")[:-1]
+            subprocess.check_output(  # noqa: DUO116
-    )
+                command, shell=True, timeout=MEDIUM_REQUEST_TIMEOUT
-    USERS.remove("Public")
+            )
-
+            .decode()
-    STARTUP_FILES_PER_USER = [
+            .split("\r\n")[:-1]
-        "\\".join(
-            SHELL_STARTUP_FILE_PATH_COMPONENTS[:2] + [user] + SHELL_STARTUP_FILE_PATH_COMPONENTS[3:]
         )
-        for user in USERS
+        users.remove("Public")
+    except subprocess.TimeoutExpired:
+        logger.error(f"Command {command} timed out")
+        return "", []
+
+    startup_files_per_user = [
+        "\\".join(
+            shell_startup_file_path_components[:2] + [user] + shell_startup_file_path_components[3:]
+        )
+        for user in users
     ]
 
     return [
         "powershell.exe",
         str(MODIFY_POWERSHELL_STARTUP_SCRIPT),
         "-startup_file_path {0}",
-    ], STARTUP_FILES_PER_USER
+    ], startup_files_per_user