From ce7362e27881a0c3043c34c02ea909d1d39cea7f Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Thu, 2 Dec 2021 14:26:10 +0200
Subject: [PATCH] Agent: add a waiting timer to allow exploited machines to
 connect to the tunnel (in agent cleanup)

---
 monkey/infection_monkey/monkey.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index cc3c5156e..4eb959129 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -3,6 +3,7 @@ import logging
 import os
 import subprocess
 import sys
+import time
 
 import infection_monkey.tunnel as tunnel
 from common.utils.attack_utils import ScanStatus, UsageEnum
@@ -186,6 +187,7 @@ class InfectionMonkey:
 
     def cleanup(self):
         logger.info("Monkey cleanup started")
+        self._wait_for_exploited_machine_connection()
         try:
             if self._is_upgrade_to_64_needed():
                 logger.debug("Cleanup not needed for 32 bit agent on 64 bit system(it didn't run)")
@@ -220,6 +222,19 @@ class InfectionMonkey:
 
         logger.info("Monkey is shutting down")
 
+    def _wait_for_exploited_machine_connection(self):
+        # TODO check for actual exploitation
+        machines_exploited = False
+        # if host was exploited, before continue to closing the tunnel ensure the exploited
+        # host had its chance to
+        # connect to the tunnel
+        if machines_exploited:
+            time_to_sleep = WormConfiguration.keep_tunnel_open_time
+            logger.info(
+                "Sleeping %d seconds for exploited machines to connect to tunnel", time_to_sleep
+            )
+            time.sleep(time_to_sleep)
+
     @staticmethod
     def _close_tunnel():
         tunnel_address = (