Agent: Add agent deletion check

This check will avoid duplicate deletion calls in case an exception happens
This commit is contained in:
vakarisz 2022-04-05 15:30:38 +03:00
parent e40703dcac
commit 64f2d598cf
1 changed files with 9 additions and 3 deletions

View File

@ -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