From 96069d3ae69a8dc35c5f7dceabcb4e8d40ea259a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 10 Mar 2022 08:32:54 -0500 Subject: [PATCH] Agent: Wrap get_log_path() with easier to use functions --- monkey/infection_monkey/main.py | 6 +++--- monkey/infection_monkey/monkey.py | 4 ++-- monkey/infection_monkey/utils/monkey_log_path.py | 9 +++++++-- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/main.py b/monkey/infection_monkey/main.py index f3e6b0a01..d6523bbcd 100644 --- a/monkey/infection_monkey/main.py +++ b/monkey/infection_monkey/main.py @@ -16,7 +16,7 @@ from infection_monkey.config import EXTERNAL_CONFIG_FILE, WormConfiguration from infection_monkey.dropper import MonkeyDrops from infection_monkey.model import DROPPER_ARG, MONKEY_ARG from infection_monkey.monkey import InfectionMonkey -from infection_monkey.utils.monkey_log_path import get_log_path +from infection_monkey.utils.monkey_log_path import get_agent_log_path, get_dropper_log_path logger = None @@ -80,10 +80,10 @@ def main(): try: if MONKEY_ARG == monkey_mode: - log_path = get_log_path("agent") + log_path = get_agent_log_path() monkey_cls = InfectionMonkey elif DROPPER_ARG == monkey_mode: - log_path = get_log_path("dropper") + log_path = get_dropper_log_path() monkey_cls = MonkeyDrops else: return True diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 0035b5cf6..a62547ebc 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -52,7 +52,7 @@ from infection_monkey.utils.monkey_dir import ( get_monkey_dir_path, remove_monkey_dir, ) -from infection_monkey.utils.monkey_log_path import get_log_path +from infection_monkey.utils.monkey_log_path import get_agent_log_path from infection_monkey.utils.signal_handler import register_signal_handlers, reset_signal_handlers logger = logging.getLogger(__name__) @@ -288,7 +288,7 @@ class InfectionMonkey: @staticmethod def _send_log(): - monkey_log_path = get_log_path("agent") + monkey_log_path = get_agent_log_path() if os.path.exists(monkey_log_path): with open(monkey_log_path, "r") as f: log = f.read() diff --git a/monkey/infection_monkey/utils/monkey_log_path.py b/monkey/infection_monkey/utils/monkey_log_path.py index bad203542..4708213fa 100644 --- a/monkey/infection_monkey/utils/monkey_log_path.py +++ b/monkey/infection_monkey/utils/monkey_log_path.py @@ -2,11 +2,12 @@ import os import sys import tempfile import time -from functools import lru_cache +from functools import lru_cache, partial +# Cache the result of the call so that subsequent calls always return the same result @lru_cache(maxsize=None) -def get_log_path(monkey_arg: str): +def _get_log_path(monkey_arg: str) -> str: return ( os.path.expandvars(_generate_random_log_filepath(monkey_arg)) if sys.platform == "win32" @@ -21,3 +22,7 @@ def _generate_random_log_filepath(monkey_arg: str) -> str: _, monkey_log_path = tempfile.mkstemp(suffix=suffix, prefix=prefix) return monkey_log_path + + +get_agent_log_path = partial(_get_log_path, "monkey") +get_dropper_log_path = partial(_get_log_path, "dropper")