From 784e3839595a6a0260efc5fc926bc979569d8e9d Mon Sep 17 00:00:00 2001 From: Itay Mizeretz Date: Wed, 28 Feb 2018 19:38:05 +0200 Subject: [PATCH] Check if should upgrade only once Don't send state-done telemetry if upgrading --- infection_monkey/monkey.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/infection_monkey/monkey.py b/infection_monkey/monkey.py index 0ca85b6b8..9e2d54fd1 100644 --- a/infection_monkey/monkey.py +++ b/infection_monkey/monkey.py @@ -36,6 +36,7 @@ class InfectionMonkey(object): self._default_server = None self._depth = 0 self._opts = None + self._upgrading_to_64 = False def initialize(self): LOG.info("Monkey is initializing...") @@ -75,6 +76,7 @@ class InfectionMonkey(object): ControlClient.find_server(default_tunnel=self._default_tunnel) if WindowsUpgrader.should_upgrade(): + self._upgrading_to_64 = True LOG.info("32bit monkey running on 64bit Windows. Upgrading.") WindowsUpgrader.upgrade(self._opts) return @@ -225,7 +227,8 @@ class InfectionMonkey(object): self._keep_running = False # Signal the server (before closing the tunnel) - ControlClient.send_telemetry("state", {'done': True}) + if not self._upgrading_to_64: + ControlClient.send_telemetry("state", {'done': True}) # Close tunnel tunnel_address = ControlClient.proxies.get('https', '').replace('https://', '').split(':')[0] @@ -235,11 +238,11 @@ class InfectionMonkey(object): firewall.close() - if not WindowsUpgrader.should_upgrade(): + if not self._upgrading_to_64: self._singleton.unlock() if WormConfiguration.self_delete_in_cleanup \ - and -1 == sys.executable.find('python') and not WindowsUpgrader.should_upgrade(): + and -1 == sys.executable.find('python') and not self._upgrading_to_64: try: if "win32" == sys.platform: from _subprocess import SW_HIDE, STARTF_USESHOWWINDOW, CREATE_NEW_CONSOLE