Agent: Add timeout to scheduling jobs PBA

monkey/infection_monkey/post_breach/job_scheduling/job_scheduling.py

@@ -1,5 +1,8 @@
+import logging
 import subprocess
+from typing import Iterable, Tuple
 
+from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
 from infection_monkey.post_breach.job_scheduling.linux_job_scheduling import (
     get_linux_commands_to_schedule_jobs,
 )
@@ -9,8 +12,10 @@ from infection_monkey.post_breach.job_scheduling.windows_job_scheduling import (
 )
 from infection_monkey.utils.environment import is_windows_os
 
+logger = logging.getLogger(__name__)
 
-def get_commands_to_schedule_jobs():
+
+def get_commands_to_schedule_jobs() -> Tuple[Iterable[str], str]:
     linux_cmds = get_linux_commands_to_schedule_jobs()
     windows_cmds = get_windows_commands_to_schedule_jobs()
     return linux_cmds, windows_cmds
@@ -18,4 +23,13 @@ def get_commands_to_schedule_jobs():
 
 def remove_scheduled_jobs():
     if is_windows_os():
-        subprocess.run(get_windows_commands_to_remove_scheduled_jobs(), shell=True)  # noqa: DUO116
+        try:
+            subprocess.run(  # noqa: DUO116
+                get_windows_commands_to_remove_scheduled_jobs(),
+                timeout=LONG_REQUEST_TIMEOUT,
+                shell=True,
+            )
+        except subprocess.CalledProcessError as err:
+            logger.error(f"An error occured removing scheduled jobs on Windows: {err}")
+        except subprocess.TimeoutExpired as err:
+            logger.error(f"A timeout occured removing scheduled jobs on Windows: {err}")

monkey/infection_monkey/post_breach/job_scheduling/linux_job_scheduling.py

@@ -1,7 +1,9 @@
+from typing import Iterable
+
 TEMP_CRON = "$HOME/monkey-schedule-jobs"
 
 
-def get_linux_commands_to_schedule_jobs():
+def get_linux_commands_to_schedule_jobs() -> Iterable[str]:
     return [
         f"touch {TEMP_CRON} &&",
         f"crontab -l > {TEMP_CRON} &&",

monkey/infection_monkey/post_breach/job_scheduling/windows_job_scheduling.py

@@ -6,9 +6,9 @@ SCHEDULED_TASK_COMMAND = r"C:\windows\system32\cmd.exe"
 # /T1053.005.md
 
 
-def get_windows_commands_to_schedule_jobs():
+def get_windows_commands_to_schedule_jobs() -> str:
     return f"schtasks /Create /SC monthly /F /TN {SCHEDULED_TASK_NAME} /TR {SCHEDULED_TASK_COMMAND}"
 
 
-def get_windows_commands_to_remove_scheduled_jobs():
+def get_windows_commands_to_remove_scheduled_jobs() -> str:
     return f"schtasks /Delete /TN {SCHEDULED_TASK_NAME} /F > nul 2>&1"