From 9ac4d23f2883aa0bc4b41d4265c869c249d4402c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 1 Apr 2022 12:04:59 +0530 Subject: [PATCH] Agent: Catch timeout error in PBA base class --- monkey/infection_monkey/post_breach/pba.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py index 0920513bd..0ef8e0ecb 100644 --- a/monkey/infection_monkey/post_breach/pba.py +++ b/monkey/infection_monkey/post_breach/pba.py @@ -19,7 +19,12 @@ class PBA: """ def __init__( - self, telemetry_messenger: ITelemetryMessenger, name="unknown", linux_cmd="", windows_cmd="", timeout: int = LONG_REQUEST_TIMEOUT + self, + telemetry_messenger: ITelemetryMessenger, + name="unknown", + linux_cmd="", + windows_cmd="", + timeout: int = LONG_REQUEST_TIMEOUT, ): """ :param name: Name of post breach action. @@ -78,9 +83,10 @@ class PBA: self.command, stderr=subprocess.STDOUT, shell=True, timeout=self.timeout ).decode() return output, True - except subprocess.CalledProcessError as e: - # Return error output of the command - return e.output.decode(), False + except subprocess.CalledProcessError as err: + return err.output.decode(), False + except subprocess.TimeoutExpired as err: + return str(err), False @staticmethod def choose_command(linux_cmd, windows_cmd):