From 20d5fb3748dad5c91230c7529f37db5e1259f59a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 13:05:22 -0400 Subject: [PATCH] Island: Use re.Pattern for IFileRepository.delete_files_by_regex() --- .../cc/repository/file_agent_log_repository.py | 15 ++++++++------- .../file_repository_caching_decorator.py | 3 ++- .../file_repository_locking_decorator.py | 3 ++- .../file_repository_logging_decorator.py | 3 ++- .../cc/repository/i_file_repository.py | 3 ++- .../tests/monkey_island/mock_file_repository.py | 3 ++- 6 files changed, 18 insertions(+), 12 deletions(-) diff --git a/monkey/monkey_island/cc/repository/file_agent_log_repository.py b/monkey/monkey_island/cc/repository/file_agent_log_repository.py index d92e07170..3ded768dd 100644 --- a/monkey/monkey_island/cc/repository/file_agent_log_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_log_repository.py @@ -9,7 +9,6 @@ from monkey_island.cc.repository import ( RetrievalError, ) -AGENT_LOG_FILE_NAME_PATTERN = "agent-*.log" AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$") @@ -18,13 +17,14 @@ class FileAgentLogRepository(IAgentLogRepository): self._file_repository = file_repository def upsert_agent_log(self, agent_id: AgentID, log_contents: str): - self._file_repository.save_file( - self._get_agent_log_file_name(agent_id), io.BytesIO(log_contents.encode()) - ) + log_file_name = FileAgentLogRepository._get_agent_log_file_name(agent_id) + self._file_repository.save_file(log_file_name, io.BytesIO(log_contents.encode())) def get_agent_log(self, agent_id: AgentID) -> str: + log_file_name = FileAgentLogRepository._get_agent_log_file_name(agent_id) + try: - with self._file_repository.open_file(self._get_agent_log_file_name(agent_id)) as f: + with self._file_repository.open_file(log_file_name) as f: log_contents = f.read().decode() return log_contents except RepositoryError as err: @@ -35,5 +35,6 @@ class FileAgentLogRepository(IAgentLogRepository): def reset(self): self._file_repository.delete_files_by_regex(AGENT_LOG_FILE_NAME_REGEX) - def _get_agent_log_file_name(self, agent_id: AgentID) -> str: - return AGENT_LOG_FILE_NAME_PATTERN.replace("*", str(agent_id)) + @staticmethod + def _get_agent_log_file_name(agent_id: AgentID) -> str: + return f"agent-{agent_id}.log" diff --git a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py index fbb26deaa..324ac0f89 100644 --- a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py @@ -1,4 +1,5 @@ import io +import re import shutil from functools import lru_cache from typing import BinaryIO @@ -36,7 +37,7 @@ class FileRepositoryCachingDecorator(IFileRepository): self._open_file.cache_clear() return self._file_repository.delete_file(unsafe_file_name) - def delete_files_by_regex(self, file_name_regex: str): + def delete_files_by_regex(self, file_name_regex: re.Pattern): self._open_file.cache_clear() return self._file_repository.delete_files_by_regex(file_name_regex) diff --git a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py index d1f772cbc..72af1206a 100644 --- a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -1,3 +1,4 @@ +import re from typing import BinaryIO from readerwriterlock import rwlock @@ -26,7 +27,7 @@ class FileRepositoryLockingDecorator(IFileRepository): with self._rwlock.gen_wlock(): return self._file_repository.delete_file(unsafe_file_name) - def delete_files_by_regex(self, file_name_regex: str): + def delete_files_by_regex(self, file_name_regex: re.Pattern): with self._rwlock.gen_wlock(): return self._file_repository.delete_files_by_regex(file_name_regex) diff --git a/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py index b4b0522b4..2ddb31fe1 100644 --- a/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py @@ -1,4 +1,5 @@ import logging +import re from typing import BinaryIO from . import IFileRepository @@ -26,7 +27,7 @@ class FileRepositoryLoggingDecorator(IFileRepository): logger.debug(f"Deleting file {unsafe_file_name}") return self._file_repository.delete_file(unsafe_file_name) - def delete_files_by_regex(self, file_name_regex: str): + def delete_files_by_regex(self, file_name_regex: re.Pattern): logger.debug(f'Deleting files whose names match the regex "{file_name_regex}"') return self._file_repository.delete_files_by_regex(file_name_regex) diff --git a/monkey/monkey_island/cc/repository/i_file_repository.py b/monkey/monkey_island/cc/repository/i_file_repository.py index 3461186d2..f28fc2bbd 100644 --- a/monkey/monkey_island/cc/repository/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/i_file_repository.py @@ -1,4 +1,5 @@ import abc +import re from typing import BinaryIO from monkey_island.cc.repository import UnknownRecordError @@ -51,7 +52,7 @@ class IFileRepository(metaclass=abc.ABCMeta): pass @abc.abstractmethod - def delete_files_by_regex(self, file_name_regex: str): + def delete_files_by_regex(self, file_name_regex: re.Pattern): """ Delete files whose names match a particular regex diff --git a/monkey/tests/monkey_island/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py index e7dc76bc1..0d6e61e9b 100644 --- a/monkey/tests/monkey_island/mock_file_repository.py +++ b/monkey/tests/monkey_island/mock_file_repository.py @@ -1,4 +1,5 @@ import io +import re from typing import BinaryIO from monkey_island.cc import repository @@ -24,7 +25,7 @@ class MockFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): pass - def delete_files_by_regex(self, file_name_regex: str): + def delete_files_by_regex(self, file_name_regex: re.Pattern): pass def delete_all_files(self):