Agent: Add timeouts to signed script PBA

This commit is contained in:
Shreya Malviya 2022-03-31 17:29:16 +05:30
parent 31ae13ed0b
commit 4cc57f1236
2 changed files with 10 additions and 3 deletions

View File

@ -3,6 +3,7 @@ import subprocess
from typing import Dict from typing import Dict
from common.common_consts.post_breach_consts import POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC 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.pba import PBA
from infection_monkey.post_breach.signed_script_proxy.signed_script_proxy import ( from infection_monkey.post_breach.signed_script_proxy.signed_script_proxy import (
cleanup_changes, cleanup_changes,
@ -21,6 +22,7 @@ class SignedScriptProxyExecution(PBA):
telemetry_messenger, telemetry_messenger,
POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC, POST_BREACH_SIGNED_SCRIPT_PROXY_EXEC,
windows_cmd=" ".join(windows_cmds), windows_cmd=" ".join(windows_cmds),
timeout=MEDIUM_REQUEST_TIMEOUT,
) )
def run(self, options: Dict): def run(self, options: Dict):
@ -28,7 +30,7 @@ class SignedScriptProxyExecution(PBA):
try: try:
if is_windows_os(): if is_windows_os():
original_comspec = subprocess.check_output( # noqa: DUO116 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() ).decode()
super().run(options) super().run(options)
return self.pba_data return self.pba_data

View File

@ -1,5 +1,6 @@
import subprocess import subprocess
from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
from infection_monkey.post_breach.signed_script_proxy.windows.signed_script_proxy import ( from infection_monkey.post_breach.signed_script_proxy.windows.signed_script_proxy import (
get_windows_commands_to_delete_temp_comspec, get_windows_commands_to_delete_temp_comspec,
get_windows_commands_to_proxy_execution_using_signed_script, 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): def cleanup_changes(original_comspec):
if is_windows_os(): if is_windows_os():
subprocess.run( # noqa: DUO116 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