From 64f2d598cf891e1335df52da366340d12c54d7d3 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Tue, 5 Apr 2022 15:30:38 +0300
Subject: [PATCH] Agent: Add agent deletion check

This check will avoid duplicate deletion calls in case an exception happens
---
 monkey/infection_monkey/monkey.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 96f1873bc..922f315a1 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -343,6 +343,7 @@ class InfectionMonkey:
 
     def cleanup(self):
         logger.info("Monkey cleanup started")
+        deleted = None
         try:
             if self._master:
                 self._master.cleanup()
@@ -357,7 +358,7 @@ class InfectionMonkey:
                 firewall.remove_firewall_rule()
                 firewall.close()
 
-            InfectionMonkey._self_delete()
+            deleted = InfectionMonkey._self_delete()
 
             InfectionMonkey._send_log()
 
@@ -371,7 +372,8 @@ class InfectionMonkey:
             self._singleton.unlock()
         except Exception as e:
             logger.error(f"An error occurred while cleaning up the monkey agent: {e}")
-            InfectionMonkey._self_delete()
+            if deleted is None:
+                InfectionMonkey._self_delete()
 
         logger.info("Monkey is shutting down")
 
@@ -400,9 +402,10 @@ class InfectionMonkey:
         ControlClient.send_log(log)
 
     @staticmethod
-    def _self_delete():
+    def _self_delete() -> bool:
         status = ScanStatus.USED if remove_monkey_dir() else ScanStatus.SCANNED
         T1107Telem(status, get_monkey_dir_path()).send()
+        deleted = False
 
         if -1 == sys.executable.find("python"):
             try:
@@ -421,11 +424,14 @@ class InfectionMonkey:
                         close_fds=True,
                         startupinfo=startupinfo,
                     )
+                    deleted = True
                 else:
                     os.remove(sys.executable)
                     status = ScanStatus.USED
+                    deleted = True
             except Exception as exc:
                 logger.error("Exception in self delete: %s", exc)
                 status = ScanStatus.SCANNED
             if status:
                 T1107Telem(status, sys.executable).send()
+            return deleted