From ca485bf569cfb6ddef053df65c59d9a3c473c86f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 7 Mar 2022 03:59:47 -0500 Subject: [PATCH] Agent: Return temporary monkey_dir as Path instead of str --- .../infection_monkey/telemetry/attack/t1107_telem.py | 2 +- monkey/infection_monkey/utils/monkey_dir.py | 12 ++++++++---- .../telemetry/attack/test_t1107_telem.py | 6 ++++++ 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/telemetry/attack/t1107_telem.py b/monkey/infection_monkey/telemetry/attack/t1107_telem.py index 816488f3b..c3667289b 100644 --- a/monkey/infection_monkey/telemetry/attack/t1107_telem.py +++ b/monkey/infection_monkey/telemetry/attack/t1107_telem.py @@ -13,5 +13,5 @@ class T1107Telem(AttackTelem): def get_data(self): data = super(T1107Telem, self).get_data() - data.update({"path": self.path}) + data.update({"path": str(self.path)}) return data diff --git a/monkey/infection_monkey/utils/monkey_dir.py b/monkey/infection_monkey/utils/monkey_dir.py index 7f74f9158..14269c7b3 100644 --- a/monkey/infection_monkey/utils/monkey_dir.py +++ b/monkey/infection_monkey/utils/monkey_dir.py @@ -1,5 +1,6 @@ import shutil import tempfile +from pathlib import Path MONKEY_DIR_PREFIX = "monkey_dir_" _monkey_dir = None @@ -7,13 +8,13 @@ _monkey_dir = None # TODO: Check if we even need this. Individual plugins can just use tempfile.mkdtemp() or # tempfile.mkftemp() if they need to. -def create_monkey_dir(): +def create_monkey_dir() -> Path: """ Creates directory for monkey and related files """ global _monkey_dir - _monkey_dir = tempfile.mkdtemp(prefix=MONKEY_DIR_PREFIX, dir=tempfile.gettempdir()) + _monkey_dir = Path(tempfile.mkdtemp(prefix=MONKEY_DIR_PREFIX, dir=tempfile.gettempdir())) return _monkey_dir @@ -29,5 +30,8 @@ def remove_monkey_dir(): return False -def get_monkey_dir_path(): - return _monkey_dir +def get_monkey_dir_path() -> Path: + if _monkey_dir is None: + create_monkey_dir() + + return _monkey_dir # type: ignore diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/attack/test_t1107_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/attack/test_t1107_telem.py index bb1bf2088..7680191a5 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/attack/test_t1107_telem.py +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/attack/test_t1107_telem.py @@ -1,4 +1,5 @@ import json +from pathlib import Path import pytest @@ -20,3 +21,8 @@ def test_T1107_send(T1107_telem_test_instance, spy_send_telemetry): expected_data = json.dumps(expected_data, cls=T1107_telem_test_instance.json_encoder) assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "attack" + + +def test_T1107_send__path(spy_send_telemetry): + T1107Telem(STATUS, Path(PATH)).send() + assert json.loads(spy_send_telemetry.data)["path"] == PATH