Agent: Wrap get_log_path() with easier to use functions

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

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()

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")