Agent: Return temporary monkey_dir as Path instead of str

This commit is contained in:
Mike Salvatore 2022-03-07 03:59:47 -05:00
parent efa0c5beb4
commit ca485bf569
3 changed files with 15 additions and 5 deletions

View File

@ -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

View File

@ -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

View File

@ -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