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 import subprocess
from typing import Dict, Iterable from typing import Dict, Iterable
from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
from common.utils.attack_utils import ScanStatus from common.utils.attack_utils import ScanStatus
from infection_monkey.i_puppet.i_puppet import PostBreachData from infection_monkey.i_puppet.i_puppet import PostBreachData
from infection_monkey.telemetry.attack.t1064_telem import T1064Telem from infection_monkey.telemetry.attack.t1064_telem import T1064Telem
@ -18,7 +19,7 @@ class PBA:
""" """
def __init__( 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. :param name: Name of post breach action.
@ -29,6 +30,7 @@ class PBA:
self.name = name self.name = name
self.pba_data = [] self.pba_data = []
self.telemetry_messenger = telemetry_messenger self.telemetry_messenger = telemetry_messenger
self.timeout = timeout
def run(self, options: Dict) -> Iterable[PostBreachData]: def run(self, options: Dict) -> Iterable[PostBreachData]:
""" """
@ -73,7 +75,7 @@ class PBA:
""" """
try: try:
output = subprocess.check_output( # noqa: DUO116 output = subprocess.check_output( # noqa: DUO116
self.command, stderr=subprocess.STDOUT, shell=True self.command, stderr=subprocess.STDOUT, shell=True, timeout=self.timeout
).decode() ).decode()
return output, True return output, True
except subprocess.CalledProcessError as e: except subprocess.CalledProcessError as e: