p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Add timeout to PBA base class's run()

This commit is contained in:
Shreya Malviya 2022-03-31 13:52:24 +05:30
parent dba9b9a637
commit 31ae13ed0b
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
monkey/infection_monkey/post_breach/pba.py
View File

@ -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,7 @@ 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 +30,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,7 +75,7 @@ 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: