diff --git a/monkey/infection_monkey/dropper.py b/monkey/infection_monkey/dropper.py index 7fdd6859c87..09a2ec7fd08 100644 --- a/monkey/infection_monkey/dropper.py +++ b/monkey/infection_monkey/dropper.py @@ -9,8 +9,9 @@ import time from pathlib import PosixPath, WindowsPath +from common import OperatingSystem from common.utils.argparse_types import positive_int -from common.utils.environment import is_windows_os +from common.utils.environment import get_os from infection_monkey.utils.commands import ( build_monkey_commandline_explicitly, get_monkey_commandline_linux, @@ -18,14 +19,6 @@ ) from infection_monkey.utils.file_utils import mark_file_for_deletion_on_windows -if "win32" == sys.platform: - from win32process import DETACHED_PROCESS - - DATE_REFERENCE_PATH_WINDOWS = os.path.expandvars(WindowsPath(r"%windir%\system32\kernel32.dll")) -else: - DETACHED_PROCESS = 0 - DATE_REFERENCE_PATH_LINUX = PosixPath("/bin/sh") - # Linux doesn't have WindowsError try: WindowsError @@ -39,6 +32,20 @@ MOVEFILE_DELAY_UNTIL_REBOOT = 4 +def file_exists_at_destination(source_path, destination_path) -> bool: + try: + return filecmp.cmp(source_path, destination_path) + except OSError: + return False + + +def get_date_reference_path(): + if get_os() == OperatingSystem.WINDOWS: + return os.path.expandvars(WindowsPath(r"%windir%\system32\kernel32.dll")) + else: + return PosixPath("/bin/sh") + + class MonkeyDrops(object): def __init__(self, args): arg_parser = argparse.ArgumentParser() @@ -61,73 +68,71 @@ def start(self): logger.error("No destination path specified") return False + source_path = self._config["source_path"] + destination_path = self._config["destination_path"] + # we copy/move only in case path is different + file_exists = file_exists_at_destination(source_path, destination_path) + if not file_exists and os.path.exists(destination_path): + os.remove(destination_path) + + if ( + not file_exists + and not self._move_file(source_path, destination_path) + and not self._copy_file(source_path, destination_path) + ): + return False + + MonkeyDrops._try_update_access_time(destination_path) + monkey_process = self._run_monkey(destination_path) + + time.sleep(3) + if monkey_process.poll() is not None: + logger.warning("Seems like monkey died too soon") + + def _move_file(self, source_path, destination_path) -> bool: try: - file_moved = filecmp.cmp(self._config["source_path"], self._config["destination_path"]) - except OSError: - file_moved = False + shutil.move(source_path, destination_path) + logger.info(f"Moved source file '{source_path}' into '{destination_path}'") + except (WindowsError, IOError, OSError) as exc: + logger.debug( + f"Error moving source file '{source_path}' into '{destination_path}': {exc}" + ) - if not file_moved and os.path.exists(self._config["destination_path"]): - os.remove(self._config["destination_path"]) + return False - # always try to move the file first - if not file_moved: - try: - shutil.move(self._config["source_path"], self._config["destination_path"]) - - logger.info( - "Moved source file '%s' into '%s'", - self._config["source_path"], - self._config["destination_path"], - ) - - file_moved = True - except (WindowsError, IOError, OSError) as exc: - logger.debug( - "Error moving source file '%s' into '%s': %s", - self._config["source_path"], - self._config["destination_path"], - exc, - ) - - # if file still need to change path, copy it - if not file_moved: - try: - shutil.copy(self._config["source_path"], self._config["destination_path"]) - - logger.info( - "Copied source file '%s' into '%s'", - self._config["source_path"], - self._config["destination_path"], - ) - except (WindowsError, IOError, OSError) as exc: - logger.error( - "Error copying source file '%s' into '%s': %s", - self._config["source_path"], - self._config["destination_path"], - exc, - ) - - return False - - if sys.platform == "win32": - dropper_date_reference_path = DATE_REFERENCE_PATH_WINDOWS - else: - dropper_date_reference_path = DATE_REFERENCE_PATH_LINUX + return True + + def _copy_file(self, source_path, destination_path) -> bool: + try: + shutil.copy(source_path, destination_path) + logger.info(f"Copied source file '{source_path}' into '{destination_path}'") + except (WindowsError, IOError, OSError) as exc: + logger.debug( + f"Error copying source file '{source_path}' into '{destination_path}': {exc}" + ) + + return False + + return True + + @staticmethod + def _try_update_access_time(destination_path): + dropper_date_reference_path = get_date_reference_path() try: ref_stat = os.stat(dropper_date_reference_path) except OSError: logger.warning( - "Cannot set reference date using '%s', file not found", - dropper_date_reference_path, + f"Cannot set reference date using '{dropper_date_reference_path}', file not found" ) else: try: - os.utime(self._config["destination_path"], (ref_stat.st_atime, ref_stat.st_mtime)) + os.utime(destination_path, (ref_stat.st_atime, ref_stat.st_mtime)) except OSError: logger.warning("Cannot set reference date to destination file") + def _run_monkey(self, destination_path) -> subprocess.Popen: monkey_options = build_monkey_commandline_explicitly( parent=self.opts.parent, servers=self.opts.servers, @@ -135,10 +140,10 @@ def start(self): location=None, ) - if is_windows_os(): - monkey_commandline = get_monkey_commandline_windows( - self._config["destination_path"], monkey_options - ) + if get_os() == OperatingSystem.WINDOWS: + from win32process import DETACHED_PROCESS + + monkey_commandline = get_monkey_commandline_windows(destination_path, monkey_options) monkey_process = subprocess.Popen( monkey_commandline, @@ -149,11 +154,10 @@ def start(self): creationflags=DETACHED_PROCESS, ) else: - dest_path = self._config["destination_path"] # In Linux, we need to change the directory first, which is done # using thw `cwd` argument in `subprocess.Popen` below - monkey_commandline = get_monkey_commandline_linux(dest_path, monkey_options) + monkey_commandline = get_monkey_commandline_linux(destination_path, monkey_options) monkey_process = subprocess.Popen( monkey_commandline, @@ -161,38 +165,34 @@ def start(self): stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True, - cwd="/".join(dest_path.split("/")[0:-1]), - creationflags=DETACHED_PROCESS, + cwd="/".join(destination_path.split("/")[0:-1]), ) logger.info( - "Executed monkey process (PID=%d) with command line: %s", - monkey_process.pid, - " ".join(monkey_commandline), + f"Executed monkey process (PID={monkey_process.pid}) " + f"with command line: {' '.join(monkey_commandline)}" ) - - time.sleep(3) - if monkey_process.poll() is not None: - logger.warning("Seems like monkey died too soon") + return monkey_process def cleanup(self): logger.info("Cleaning up the dropper") + source_path = self._config["source_path"] + try: - if self._config["source_path"].lower() != self._config[ - "destination_path" - ].lower() and os.path.exists(self._config["source_path"]): - - # try removing the file first - try: - os.remove(self._config["source_path"]) - except Exception as exc: - logger.debug( - "Error removing source file '%s': %s", self._config["source_path"], exc - ) - - # mark the file for removal on next boot - mark_file_for_deletion_on_windows(WindowsPath(self._config["source_path"])) + if source_path.lower() != self._config["destination_path"].lower() and os.path.exists( + source_path + ): + self._remove_file(source_path) logger.info("Dropper cleanup complete") except AttributeError: logger.error("Invalid configuration options. Failing") + + def _remove_file(self, path): + try: + os.remove(path) + except Exception as exc: + logger.debug(f"Error removing source file '{path}': {exc}") + + # mark the file for removal on next boot + mark_file_for_deletion_on_windows(WindowsPath(path))