diff --git a/infection_monkey/utils.py b/infection_monkey/utils.py index d95407341..baef2372a 100644 --- a/infection_monkey/utils.py +++ b/infection_monkey/utils.py @@ -1,5 +1,6 @@ import os import sys +import struct from config import WormConfiguration @@ -12,3 +13,15 @@ def get_monkey_log_path(): def get_dropper_log_path(): return os.path.expandvars(WormConfiguration.dropper_log_path_windows) if sys.platform == "win32" \ else WormConfiguration.dropper_log_path_linux + + +def is_64bit_os(): + return 'PROGRAMFILES(X86)' in os.environ + + +def is_64bit_python(): + return struct.calcsize("P") == 8 + + +def is_windows_os(): + return sys.platform.startswith("win") diff --git a/infection_monkey/windows_upgrader.py b/infection_monkey/windows_upgrader.py index 996e2a856..38cb3a479 100644 --- a/infection_monkey/windows_upgrader.py +++ b/infection_monkey/windows_upgrader.py @@ -1,6 +1,4 @@ import logging -import os -import struct import subprocess import sys import shutil @@ -12,6 +10,7 @@ from config import WormConfiguration from control import ControlClient from exploit.tools import build_monkey_commandline_explicitly from model import MONKEY_CMDLINE_WINDOWS +from utils import is_windows_os, is_64bit_os, is_64bit_python __author__ = 'itay.mizeretz' @@ -26,22 +25,10 @@ else: class WindowsUpgrader(object): __UPGRADE_WAIT_TIME__ = 3 - @staticmethod - def is_64bit_os(): - return 'PROGRAMFILES(X86)' in os.environ - - @staticmethod - def is_64bit_python(): - return struct.calcsize("P") == 8 - - @staticmethod - def is_windows_os(): - return sys.platform.startswith("win") - @staticmethod def should_upgrade(): - return WindowsUpgrader.is_windows_os() and WindowsUpgrader.is_64bit_os() \ - and not WindowsUpgrader.is_64bit_python() + return is_windows_os() and is_64bit_os() \ + and not is_64bit_python() @staticmethod def upgrade(opts):