Check if should upgrade only once

Don't send state-done telemetry if upgrading
This commit is contained in:
Itay Mizeretz 2018-02-28 19:38:05 +02:00
parent abd738acbc
commit 784e383959
1 changed files with 6 additions and 3 deletions

View File

@ -36,6 +36,7 @@ class InfectionMonkey(object):
self._default_server = None self._default_server = None
self._depth = 0 self._depth = 0
self._opts = None self._opts = None
self._upgrading_to_64 = False
def initialize(self): def initialize(self):
LOG.info("Monkey is initializing...") LOG.info("Monkey is initializing...")
@ -75,6 +76,7 @@ class InfectionMonkey(object):
ControlClient.find_server(default_tunnel=self._default_tunnel) ControlClient.find_server(default_tunnel=self._default_tunnel)
if WindowsUpgrader.should_upgrade(): if WindowsUpgrader.should_upgrade():
self._upgrading_to_64 = True
LOG.info("32bit monkey running on 64bit Windows. Upgrading.") LOG.info("32bit monkey running on 64bit Windows. Upgrading.")
WindowsUpgrader.upgrade(self._opts) WindowsUpgrader.upgrade(self._opts)
return return
@ -225,7 +227,8 @@ class InfectionMonkey(object):
self._keep_running = False self._keep_running = False
# Signal the server (before closing the tunnel) # 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 # Close tunnel
tunnel_address = ControlClient.proxies.get('https', '').replace('https://', '').split(':')[0] tunnel_address = ControlClient.proxies.get('https', '').replace('https://', '').split(':')[0]
@ -235,11 +238,11 @@ class InfectionMonkey(object):
firewall.close() firewall.close()
if not WindowsUpgrader.should_upgrade(): if not self._upgrading_to_64:
self._singleton.unlock() self._singleton.unlock()
if WormConfiguration.self_delete_in_cleanup \ 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: try:
if "win32" == sys.platform: if "win32" == sys.platform:
from _subprocess import SW_HIDE, STARTF_USESHOWWINDOW, CREATE_NEW_CONSOLE from _subprocess import SW_HIDE, STARTF_USESHOWWINDOW, CREATE_NEW_CONSOLE