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

View File

@ -1,5 +1,6 @@
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,
@ -16,6 +17,10 @@ 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
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