From d9b55a5c214c30a5d4d846bf06cc394e6ae2e590 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 15:35:58 +0530 Subject: [PATCH 01/28] Island: Add delete_files_by_pattern() to IFileRepository --- .../monkey_island/cc/repository/i_file_repository.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/monkey/monkey_island/cc/repository/i_file_repository.py b/monkey/monkey_island/cc/repository/i_file_repository.py index bedc9cfdb..5d44a8b4a 100644 --- a/monkey/monkey_island/cc/repository/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/i_file_repository.py @@ -49,6 +49,17 @@ class IFileRepository(metaclass=abc.ABCMeta): """ pass + @abc.abstractmethod + def delete_files_by_pattern(self, file_name_pattern: str): + """ + Delete files whose names match a particular pattern + + This method matches relevant files and deletes them using `delete_file()`. + + :param file_name_pattern: A file name pattern that should be matched to delete files + """ + pass + @abc.abstractmethod def delete_all_files(self): """ From ce0affb1ed76d05408b7f71055f1026f98700c37 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 15:39:14 +0530 Subject: [PATCH 02/28] Island: Implement delete_files_by_pattern() in LocalStorageFileRepository --- .../cc/repository/local_storage_file_repository.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/monkey/monkey_island/cc/repository/local_storage_file_repository.py b/monkey/monkey_island/cc/repository/local_storage_file_repository.py index 58fd20bdb..840ab1701 100644 --- a/monkey/monkey_island/cc/repository/local_storage_file_repository.py +++ b/monkey/monkey_island/cc/repository/local_storage_file_repository.py @@ -1,3 +1,4 @@ +import glob import logging import shutil from pathlib import Path @@ -55,6 +56,10 @@ class LocalStorageFileRepository(IFileRepository): f'Error retrieving file "{unsafe_file_name}" from the repository: {err}' ) + def delete_files_by_pattern(self, file_name_pattern: str): + for file_name in glob.iglob(f"{self._storage_directory}/{file_name_pattern}"): + self.delete_file(file_name) + def delete_file(self, unsafe_file_name: str): try: safe_file_path = self._get_safe_file_path(unsafe_file_name) From 1c6cfa1ce6ac9fa2dfc49d0e1b9d2a66797821eb Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 15:39:57 +0530 Subject: [PATCH 03/28] Island: Add FileAgentLogRepository --- .../repository/file_agent_log_repository.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 monkey/monkey_island/cc/repository/file_agent_log_repository.py diff --git a/monkey/monkey_island/cc/repository/file_agent_log_repository.py b/monkey/monkey_island/cc/repository/file_agent_log_repository.py new file mode 100644 index 000000000..7abdab992 --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_agent_log_repository.py @@ -0,0 +1,35 @@ +import io + +from monkey_island.cc import repository +from monkey_island.cc.models import AgentID +from monkey_island.cc.repository import ( + IAgentLogRepository, + IFileRepository, + RetrievalError, + UnknownRecordError, +) + +AGENT_LOG_FILE_NAME_PREFIX = "agent_log_" + + +class FileAgentLogRepository(IAgentLogRepository): + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + + def upsert_agent_log(self, agent_id: AgentID, log_contents: str): + self._file_repository.save_file( + f"{AGENT_LOG_FILE_NAME_PREFIX}{agent_id}", io.BytesIO(log_contents.encode()) + ) + + def get_agent_log(self, agent_id: AgentID) -> str: + try: + with self._file_repository.open_file(f"{AGENT_LOG_FILE_NAME_PREFIX}{agent_id}") as f: + log_contents = f.read().decode() + return log_contents + except repository.FileNotFoundError as err: + raise UnknownRecordError(err) + except Exception as err: + raise RetrievalError(f"Error retrieving the agent logs: {err}") + + def reset(self): + self._file_repository.delete_files_by_pattern(f"{AGENT_LOG_FILE_NAME_PREFIX}*") From 01d8875f22a18a533593b1f640d6c07241fcbd38 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 15:42:07 +0530 Subject: [PATCH 04/28] Island: Register FileAgentLogRegister in DI container --- monkey/monkey_island/cc/services/initialize.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index c0dc600c3..1cd020fd8 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -21,6 +21,7 @@ from monkey_island.cc.event_queue import IIslandEventQueue, PyPubSubIslandEventQ from monkey_island.cc.repository import ( AgentBinaryRepository, FileAgentConfigurationRepository, + FileAgentLogRepository, FileRepositoryCachingDecorator, FileRepositoryLockingDecorator, FileRepositoryLoggingDecorator, @@ -28,6 +29,7 @@ from monkey_island.cc.repository import ( IAgentBinaryRepository, IAgentConfigurationRepository, IAgentEventRepository, + IAgentLogRepository, IAgentRepository, ICredentialsRepository, IFileRepository, @@ -117,6 +119,7 @@ def _register_repositories(container: DIContainer, data_dir: Path): container.register_instance(INodeRepository, container.resolve(MongoNodeRepository)) container.register_instance(IMachineRepository, _build_machine_repository(container)) container.register_instance(IAgentRepository, container.resolve(MongoAgentRepository)) + container.register_instance(IAgentLogRepository, container.resolve(FileAgentLogRepository)) def _decorate_file_repository(file_repository: IFileRepository) -> IFileRepository: From e1d139fde4407d2c1ff95db6df4c2e6bea2b70dd Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:03:08 +0530 Subject: [PATCH 05/28] Island: Add FileAgentLogRepository to cc/repository/__init__.py --- monkey/monkey_island/cc/repository/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 197cae607..91d6bed30 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -28,5 +28,6 @@ from .mongo_machine_repository import MongoMachineRepository from .mongo_agent_repository import MongoAgentRepository from .mongo_node_repository import MongoNodeRepository from .mongo_agent_event_repository import MongoAgentEventRepository +from .file_agent_log_repository import FileAgentLogRepository from .utils import initialize_machine_repository From c1a4641ffed68641d78ec40ad41a90053fa078e2 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:10:30 +0530 Subject: [PATCH 06/28] UT: Add delete_files_by_pattern() to SingleFileRepository and MockFileRepository --- monkey/tests/monkey_island/mock_file_repository.py | 3 +++ monkey/tests/monkey_island/single_file_repository.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/monkey/tests/monkey_island/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py index 782c9838b..f828cdd5a 100644 --- a/monkey/tests/monkey_island/mock_file_repository.py +++ b/monkey/tests/monkey_island/mock_file_repository.py @@ -24,5 +24,8 @@ class MockFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): pass + def delete_files_by_pattern(self, file_name_pattern: str): + pass + def delete_all_files(self): pass diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py index 462969acb..66086f0f2 100644 --- a/monkey/tests/monkey_island/single_file_repository.py +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -1,4 +1,5 @@ import io +import re from typing import BinaryIO from monkey_island.cc import repository @@ -8,9 +9,11 @@ from monkey_island.cc.repository import IFileRepository class SingleFileRepository(IFileRepository): def __init__(self): self._file = None + self._file_name = "" def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): self._file = io.BytesIO(file_contents.read()) + self._file_name = unsafe_file_name def open_file(self, unsafe_file_name: str) -> BinaryIO: if self._file is None: @@ -19,6 +22,11 @@ class SingleFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): self._file = None + self._file_name = "" + + def delete_files_by_pattern(self, file_name_pattern: str): + if re.match(file_name_pattern, self._file_name): + self.delete_file("") def delete_all_files(self): self.delete_file("") From 6aae63f9fc3ccab7655f938aef018d5e2db75798 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:10:55 +0530 Subject: [PATCH 07/28] UT: Add tests for FileAgentLogRepository --- .../test_file_agent_log_repository.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py new file mode 100644 index 000000000..035903e18 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -0,0 +1,39 @@ +from uuid import UUID + +import pytest +from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository + +from monkey_island.cc.repository import FileAgentLogRepository, RetrievalError, UnknownRecordError + +LOG_CONTENTS = "lots of useful information" +AGENT_ID = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") + + +@pytest.fixture +def repository(): + return FileAgentLogRepository(SingleFileRepository()) + + +def test_store_agent_log(repository): + repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) + retrieved_log_contents = repository.get_agent_log(AGENT_ID) + + assert retrieved_log_contents == LOG_CONTENTS + + +def test_get_agent_log__unknown_record_error(repository): + with pytest.raises(UnknownRecordError): + repository.get_agent_log(AGENT_ID) + + +def test_get_agent_log__retrieval_error(): + repository = FileAgentLogRepository(OpenErrorFileRepository()) + with pytest.raises(RetrievalError): + repository.get_agent_log(AGENT_ID) + + +def test_reset_agent_logs(repository): + repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) + repository.reset() + with pytest.raises(UnknownRecordError): + repository.get_agent_log(AGENT_ID) From 885a90728720add3aecaedb92b0b94b369c38dc7 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:21:09 +0530 Subject: [PATCH 08/28] UT: Add tests for delete_files_by_pattern() in LocalStorageFileRepository --- .../repository/test_local_storage_file_repository.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index 64541bcbe..c03d615f9 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py @@ -143,3 +143,14 @@ def test_open_locked_file(tmp_path, monkeypatch): with patch("builtins.open", Mock(side_effect=Exception())): with pytest.raises(repository.RetrievalError): fss.open_file("locked_file.txt") + + +def test_delete_files_by_pattern(tmp_path): + for filename in ["xyz-1.txt", "abc-2.txt", "pqr-3.txt", "abc-4.txt"]: + (tmp_path / filename).touch() + + fss = LocalStorageFileRepository(tmp_path) + fss.delete_files_by_pattern("abc-*") + + files = list(tmp_path.iterdir()) + assert len(files) == 2 From d5b62651a08cc78fd2539df27e676d28f2afcb98 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:23:47 +0530 Subject: [PATCH 09/28] Island: Implement delete_files_by_pattern() in FileRepositoryLoggingDecorator --- .../cc/repository/file_repository_logging_decorator.py | 4 ++++ 1 file changed, 4 insertions(+) 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 2bf8bfe6e..099b79106 100644 --- a/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py @@ -26,6 +26,10 @@ class FileRepositoryLoggingDecorator(IFileRepository): logger.debug(f"Deleting file {unsafe_file_name}") return self._file_repository.delete_file(unsafe_file_name) + def delete_files_by_pattern(self, file_name_pattern: str): + logger.debug(f'Deleting files whose names match the pattern "{file_name_pattern}"') + return self._file_repository.delete_files_by_pattern(file_name_pattern) + def delete_all_files(self): logger.debug("Deleting all files in the repository") return self._file_repository.delete_all_files() From 5f11008b40a9a1a27455b7aa4cec900ce673968e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:25:13 +0530 Subject: [PATCH 10/28] Island: Implement delete_files_by_pattern() in FileRepositoryLockingDecorator --- .../cc/repository/file_repository_locking_decorator.py | 4 ++++ 1 file changed, 4 insertions(+) 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 1bd76b4dd..602d0122a 100644 --- a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -26,6 +26,10 @@ class FileRepositoryLockingDecorator(IFileRepository): with self._rwlock.gen_wlock(): return self._file_repository.delete_file(unsafe_file_name) + def delete_files_by_pattern(self, file_name_pattern: str): + with self._rwlock.gen_wlock(): + return self._file_repository.delete_files_by_pattern(file_name_pattern) + def delete_all_files(self): with self._rwlock.gen_wlock(): return self._file_repository.delete_all_files() From cadf0d61d00c89f5c60154f21cd9bbf8078aa965 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 16:26:37 +0530 Subject: [PATCH 11/28] Island: Implement delete_files_by_pattern() in FileRepositoryCachingDecorator --- .../cc/repository/file_repository_caching_decorator.py | 4 ++++ 1 file changed, 4 insertions(+) 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 0e026d035..07bee5258 100644 --- a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py @@ -36,6 +36,10 @@ class FileRepositoryCachingDecorator(IFileRepository): self._open_file.cache_clear() return self._file_repository.delete_file(unsafe_file_name) + def delete_files_by_pattern(self, file_name_pattern: str): + self._open_file.cache_clear() + return self._file_repository.delete_files_by_pattern(file_name_pattern) + def delete_all_files(self): self._open_file.cache_clear() return self._file_repository.delete_all_files() From f61e734d29a9fc755cd97ee034606b6724e46d77 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 17:53:30 +0530 Subject: [PATCH 12/28] Island: Subscribe IAgentLogRepository.reset() to IslandEventTopic.CLEAR_SIMULATION_DATA --- monkey/monkey_island/cc/setup/island_event_handlers.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 96d2e1941..bf07be1f0 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -10,6 +10,7 @@ from monkey_island.cc.island_event_handlers import ( ) from monkey_island.cc.repository import ( IAgentEventRepository, + IAgentLogRepository, IAgentRepository, ICredentialsRepository, INodeRepository, @@ -59,9 +60,10 @@ def _subscribe_clear_simulation_data_events( island_event_queue.subscribe(topic, container.resolve(reset_machine_repository)) for i_repository in [ - INodeRepository, IAgentEventRepository, + IAgentLogRepository, IAgentRepository, + INodeRepository, ]: repository = container.resolve(i_repository) island_event_queue.subscribe(topic, repository.reset) From 30d3124cb46640c43aaa9e9e855f45c26d653837 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 18:14:48 +0530 Subject: [PATCH 13/28] UT: Fix assert logic in test_delete_files_by_pattern() --- .../cc/repository/test_local_storage_file_repository.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index c03d615f9..03fa1f5ca 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py @@ -146,11 +146,11 @@ def test_open_locked_file(tmp_path, monkeypatch): def test_delete_files_by_pattern(tmp_path): - for filename in ["xyz-1.txt", "abc-2.txt", "pqr-3.txt", "abc-4.txt"]: + for filename in {"xyz-1.txt", "abc-2.txt", "pqr-3.txt", "abc-4.txt", "abc-5.pdf"}: (tmp_path / filename).touch() fss = LocalStorageFileRepository(tmp_path) - fss.delete_files_by_pattern("abc-*") + fss.delete_files_by_pattern("abc-*.txt") - files = list(tmp_path.iterdir()) - assert len(files) == 2 + files = {f.name for f in tmp_path.iterdir()} + assert files == {"xyz-1.txt", "pqr-3.txt", "abc-5.pdf"} From c5d26749b751e745cb75fee99f9de4432d1c0fbd Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 18:29:53 +0530 Subject: [PATCH 14/28] Island: Change agent log file name in FileAgentLogRepository --- .../cc/repository/file_agent_log_repository.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 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 7abdab992..7a3cad245 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,7 @@ from monkey_island.cc.repository import ( UnknownRecordError, ) -AGENT_LOG_FILE_NAME_PREFIX = "agent_log_" +AGENT_LOG_FILE_NAME_PATTERN = "agent-*.log" class FileAgentLogRepository(IAgentLogRepository): @@ -18,12 +18,14 @@ class FileAgentLogRepository(IAgentLogRepository): def upsert_agent_log(self, agent_id: AgentID, log_contents: str): self._file_repository.save_file( - f"{AGENT_LOG_FILE_NAME_PREFIX}{agent_id}", io.BytesIO(log_contents.encode()) + AGENT_LOG_FILE_NAME_PATTERN.replace("*", agent_id), io.BytesIO(log_contents.encode()) ) def get_agent_log(self, agent_id: AgentID) -> str: try: - with self._file_repository.open_file(f"{AGENT_LOG_FILE_NAME_PREFIX}{agent_id}") as f: + with self._file_repository.open_file( + AGENT_LOG_FILE_NAME_PATTERN.replace("*", agent_id) + ) as f: log_contents = f.read().decode() return log_contents except repository.FileNotFoundError as err: @@ -32,4 +34,4 @@ class FileAgentLogRepository(IAgentLogRepository): raise RetrievalError(f"Error retrieving the agent logs: {err}") def reset(self): - self._file_repository.delete_files_by_pattern(f"{AGENT_LOG_FILE_NAME_PREFIX}*") + self._file_repository.delete_files_by_pattern(f"{AGENT_LOG_FILE_NAME_PATTERN}") From 1c486c65717307293b0270fe49fd8f5a96d56f14 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 19:27:25 +0530 Subject: [PATCH 15/28] Island: Modify delete_files_by_pattern -> delete_files_by_regex in LocalStorageFileRepository --- .../cc/repository/local_storage_file_repository.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/repository/local_storage_file_repository.py b/monkey/monkey_island/cc/repository/local_storage_file_repository.py index 840ab1701..7f7be398c 100644 --- a/monkey/monkey_island/cc/repository/local_storage_file_repository.py +++ b/monkey/monkey_island/cc/repository/local_storage_file_repository.py @@ -1,5 +1,6 @@ -import glob import logging +import os +import re import shutil from pathlib import Path from typing import BinaryIO @@ -56,9 +57,10 @@ class LocalStorageFileRepository(IFileRepository): f'Error retrieving file "{unsafe_file_name}" from the repository: {err}' ) - def delete_files_by_pattern(self, file_name_pattern: str): - for file_name in glob.iglob(f"{self._storage_directory}/{file_name_pattern}"): - self.delete_file(file_name) + def delete_files_by_regex(self, file_name_regex: re.Pattern): + for file_name in os.listdir(self._storage_directory): + if re.match(file_name_regex, file_name): + self.delete_file(file_name) def delete_file(self, unsafe_file_name: str): try: From 2707605622b32455d84e15690e594a4ff09116f6 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 19:31:12 +0530 Subject: [PATCH 16/28] Island: Add AGENT_LOG_FILE_NAME_REGEX and modify some logic in FileAgentLogRepository --- .../cc/repository/file_agent_log_repository.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 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 7a3cad245..51808cde0 100644 --- a/monkey/monkey_island/cc/repository/file_agent_log_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_log_repository.py @@ -1,4 +1,5 @@ import io +import re from monkey_island.cc import repository from monkey_island.cc.models import AgentID @@ -10,6 +11,7 @@ from monkey_island.cc.repository import ( ) AGENT_LOG_FILE_NAME_PATTERN = "agent-*.log" +AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$") class FileAgentLogRepository(IAgentLogRepository): @@ -18,14 +20,12 @@ class FileAgentLogRepository(IAgentLogRepository): def upsert_agent_log(self, agent_id: AgentID, log_contents: str): self._file_repository.save_file( - AGENT_LOG_FILE_NAME_PATTERN.replace("*", agent_id), io.BytesIO(log_contents.encode()) + self._get_agent_log_file_name(agent_id), io.BytesIO(log_contents.encode()) ) def get_agent_log(self, agent_id: AgentID) -> str: try: - with self._file_repository.open_file( - AGENT_LOG_FILE_NAME_PATTERN.replace("*", agent_id) - ) as f: + with self._file_repository.open_file(self._get_agent_log_file_name(agent_id)) as f: log_contents = f.read().decode() return log_contents except repository.FileNotFoundError as err: @@ -34,4 +34,7 @@ class FileAgentLogRepository(IAgentLogRepository): raise RetrievalError(f"Error retrieving the agent logs: {err}") def reset(self): - self._file_repository.delete_files_by_pattern(f"{AGENT_LOG_FILE_NAME_PATTERN}") + 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)) From 7823759cf8432d6646abf39673b48e7607fac920 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 19:38:01 +0530 Subject: [PATCH 17/28] Island: delete_files_by_pattern -> delete_files_by_regex --- .../cc/repository/file_repository_caching_decorator.py | 4 ++-- .../cc/repository/file_repository_locking_decorator.py | 4 ++-- .../cc/repository/file_repository_logging_decorator.py | 6 +++--- monkey/monkey_island/cc/repository/i_file_repository.py | 6 +++--- 4 files changed, 10 insertions(+), 10 deletions(-) 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 07bee5258..fbb26deaa 100644 --- a/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_caching_decorator.py @@ -36,9 +36,9 @@ class FileRepositoryCachingDecorator(IFileRepository): self._open_file.cache_clear() return self._file_repository.delete_file(unsafe_file_name) - def delete_files_by_pattern(self, file_name_pattern: str): + def delete_files_by_regex(self, file_name_regex: str): self._open_file.cache_clear() - return self._file_repository.delete_files_by_pattern(file_name_pattern) + return self._file_repository.delete_files_by_regex(file_name_regex) def delete_all_files(self): self._open_file.cache_clear() 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 602d0122a..d1f772cbc 100644 --- a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -26,9 +26,9 @@ class FileRepositoryLockingDecorator(IFileRepository): with self._rwlock.gen_wlock(): return self._file_repository.delete_file(unsafe_file_name) - def delete_files_by_pattern(self, file_name_pattern: str): + def delete_files_by_regex(self, file_name_regex: str): with self._rwlock.gen_wlock(): - return self._file_repository.delete_files_by_pattern(file_name_pattern) + return self._file_repository.delete_files_by_regex(file_name_regex) def delete_all_files(self): with self._rwlock.gen_wlock(): 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 099b79106..b4b0522b4 100644 --- a/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_logging_decorator.py @@ -26,9 +26,9 @@ class FileRepositoryLoggingDecorator(IFileRepository): logger.debug(f"Deleting file {unsafe_file_name}") return self._file_repository.delete_file(unsafe_file_name) - def delete_files_by_pattern(self, file_name_pattern: str): - logger.debug(f'Deleting files whose names match the pattern "{file_name_pattern}"') - return self._file_repository.delete_files_by_pattern(file_name_pattern) + def delete_files_by_regex(self, file_name_regex: str): + logger.debug(f'Deleting files whose names match the regex "{file_name_regex}"') + return self._file_repository.delete_files_by_regex(file_name_regex) def delete_all_files(self): logger.debug("Deleting all files in the repository") diff --git a/monkey/monkey_island/cc/repository/i_file_repository.py b/monkey/monkey_island/cc/repository/i_file_repository.py index 5d44a8b4a..877079dd3 100644 --- a/monkey/monkey_island/cc/repository/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/i_file_repository.py @@ -50,13 +50,13 @@ class IFileRepository(metaclass=abc.ABCMeta): pass @abc.abstractmethod - def delete_files_by_pattern(self, file_name_pattern: str): + def delete_files_by_regex(self, file_name_regex: str): """ - Delete files whose names match a particular pattern + Delete files whose names match a particular regex This method matches relevant files and deletes them using `delete_file()`. - :param file_name_pattern: A file name pattern that should be matched to delete files + :param file_name_regex: A regex with which a file's name should match before deleting it """ pass From 64990eea0e04f4dbe077e1c24c0cef187bd1d906 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 19:39:16 +0530 Subject: [PATCH 18/28] UT: Update tests as per changes to file repositories --- monkey/tests/monkey_island/mock_file_repository.py | 2 +- monkey/tests/monkey_island/single_file_repository.py | 4 ++-- .../cc/repository/test_local_storage_file_repository.py | 5 +++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/monkey/tests/monkey_island/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py index f828cdd5a..e7dc76bc1 100644 --- a/monkey/tests/monkey_island/mock_file_repository.py +++ b/monkey/tests/monkey_island/mock_file_repository.py @@ -24,7 +24,7 @@ class MockFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): pass - def delete_files_by_pattern(self, file_name_pattern: str): + def delete_files_by_regex(self, file_name_regex: str): pass def delete_all_files(self): diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py index 66086f0f2..be9f896ba 100644 --- a/monkey/tests/monkey_island/single_file_repository.py +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -24,8 +24,8 @@ class SingleFileRepository(IFileRepository): self._file = None self._file_name = "" - def delete_files_by_pattern(self, file_name_pattern: str): - if re.match(file_name_pattern, self._file_name): + def delete_files_by_regex(self, file_name_regex: re.Pattern): + if re.match(file_name_regex, self._file_name): self.delete_file("") def delete_all_files(self): diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py index 03fa1f5ca..f088ee16e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_local_storage_file_repository.py @@ -1,4 +1,5 @@ import io +import re from pathlib import Path from unittest.mock import Mock, patch @@ -145,12 +146,12 @@ def test_open_locked_file(tmp_path, monkeypatch): fss.open_file("locked_file.txt") -def test_delete_files_by_pattern(tmp_path): +def test_delete_files_by_regex(tmp_path): for filename in {"xyz-1.txt", "abc-2.txt", "pqr-3.txt", "abc-4.txt", "abc-5.pdf"}: (tmp_path / filename).touch() fss = LocalStorageFileRepository(tmp_path) - fss.delete_files_by_pattern("abc-*.txt") + fss.delete_files_by_regex(re.compile(r"^abc-[\w-]+.txt$")) files = {f.name for f in tmp_path.iterdir()} assert files == {"xyz-1.txt", "pqr-3.txt", "abc-5.pdf"} From 8f46b3b9fdce4d1efd3bf4b27a3901f11d87da0a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 19:55:52 +0530 Subject: [PATCH 19/28] Island: Make FileNotFoundError inherit UnknownRecordError instead of RetrievalError + add TODO for updating the rest of the code --- .../cc/repository/file_agent_log_repository.py | 10 +--------- .../monkey_island/cc/repository/i_file_repository.py | 5 +++-- 2 files changed, 4 insertions(+), 11 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 51808cde0..728d16da8 100644 --- a/monkey/monkey_island/cc/repository/file_agent_log_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_log_repository.py @@ -1,14 +1,8 @@ import io import re -from monkey_island.cc import repository from monkey_island.cc.models import AgentID -from monkey_island.cc.repository import ( - IAgentLogRepository, - IFileRepository, - RetrievalError, - UnknownRecordError, -) +from monkey_island.cc.repository import IAgentLogRepository, IFileRepository, RetrievalError AGENT_LOG_FILE_NAME_PATTERN = "agent-*.log" AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$") @@ -28,8 +22,6 @@ class FileAgentLogRepository(IAgentLogRepository): with self._file_repository.open_file(self._get_agent_log_file_name(agent_id)) as f: log_contents = f.read().decode() return log_contents - except repository.FileNotFoundError as err: - raise UnknownRecordError(err) except Exception as err: raise RetrievalError(f"Error retrieving the agent logs: {err}") diff --git a/monkey/monkey_island/cc/repository/i_file_repository.py b/monkey/monkey_island/cc/repository/i_file_repository.py index 877079dd3..3461186d2 100644 --- a/monkey/monkey_island/cc/repository/i_file_repository.py +++ b/monkey/monkey_island/cc/repository/i_file_repository.py @@ -1,10 +1,11 @@ import abc from typing import BinaryIO -from monkey_island.cc.repository import RetrievalError +from monkey_island.cc.repository import UnknownRecordError -class FileNotFoundError(RetrievalError): +# TODO: Remove this and use UnknownRecordError directly wherever needed. +class FileNotFoundError(UnknownRecordError): pass From e374341ce1804f1a3873674cc13fa675c4513879 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 28 Sep 2022 20:01:13 +0530 Subject: [PATCH 20/28] UT: Fix tests in test_file_agent_log_repository.py --- .../cc/repository/test_file_agent_log_repository.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py index 035903e18..d49d0957d 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -3,7 +3,7 @@ from uuid import UUID import pytest from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from monkey_island.cc.repository import FileAgentLogRepository, RetrievalError, UnknownRecordError +from monkey_island.cc.repository import FileAgentLogRepository, RetrievalError LOG_CONTENTS = "lots of useful information" AGENT_ID = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") @@ -22,7 +22,7 @@ def test_store_agent_log(repository): def test_get_agent_log__unknown_record_error(repository): - with pytest.raises(UnknownRecordError): + with pytest.raises(RetrievalError): repository.get_agent_log(AGENT_ID) @@ -35,5 +35,5 @@ def test_get_agent_log__retrieval_error(): def test_reset_agent_logs(repository): repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) repository.reset() - with pytest.raises(UnknownRecordError): + with pytest.raises(RetrievalError): repository.get_agent_log(AGENT_ID) From 3c2ee32bdfba2402c52325990eb7c851497ed982 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 12:21:41 -0400 Subject: [PATCH 21/28] Island: Add RepositoryError --- monkey/monkey_island/cc/repository/__init__.py | 2 +- monkey/monkey_island/cc/repository/errors.py | 14 ++++++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 91d6bed30..9012afb8b 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,4 +1,4 @@ -from .errors import RemovalError, RetrievalError, StorageError, UnknownRecordError +from .errors import RemovalError, RepositoryError, RetrievalError, StorageError, UnknownRecordError from .i_file_repository import FileNotFoundError, IFileRepository diff --git a/monkey/monkey_island/cc/repository/errors.py b/monkey/monkey_island/cc/repository/errors.py index a5c26fe15..f01114cc0 100644 --- a/monkey/monkey_island/cc/repository/errors.py +++ b/monkey/monkey_island/cc/repository/errors.py @@ -1,22 +1,28 @@ -class RemovalError(RuntimeError): +class RepositoryError(RuntimeError): + """ + Raised when a repository encounters an error while attempting any operation. + """ + + +class RemovalError(RepositoryError): """ Raised when a repository encounters an error while attempting to remove data. """ -class RetrievalError(RuntimeError): +class RetrievalError(RepositoryError): """ Raised when a repository encounters an error while attempting to retrieve data. """ -class StorageError(RuntimeError): +class StorageError(RepositoryError): """ Raised when a repository encounters an error while attempting to store data. """ -class UnknownRecordError(RuntimeError): +class UnknownRecordError(RepositoryError): """ Raised when the repository does not contain any data matching the request. """ From d49d16bc37ff343fb17327fa4bb4d72a8eddd27c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 12:38:56 -0400 Subject: [PATCH 22/28] Island: Allow RepositoryErrors to be reraised Previously, FileAgentLogRepository.get_agent_log() wrapped all errors as RetrievalError, which is not necessarily correct. This commit allows all repository errors raised by IFileRepository to be reraised, and all other, unexpected errors to be reraised as RetrievalError. --- .../repository/file_agent_log_repository.py | 9 +++++++- .../test_file_agent_log_repository.py | 22 ++++++++++++++++--- 2 files changed, 27 insertions(+), 4 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 728d16da8..d92e07170 100644 --- a/monkey/monkey_island/cc/repository/file_agent_log_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_log_repository.py @@ -2,7 +2,12 @@ import io import re from monkey_island.cc.models import AgentID -from monkey_island.cc.repository import IAgentLogRepository, IFileRepository, RetrievalError +from monkey_island.cc.repository import ( + IAgentLogRepository, + IFileRepository, + RepositoryError, + RetrievalError, +) AGENT_LOG_FILE_NAME_PATTERN = "agent-*.log" AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$") @@ -22,6 +27,8 @@ class FileAgentLogRepository(IAgentLogRepository): with self._file_repository.open_file(self._get_agent_log_file_name(agent_id)) as f: log_contents = f.read().decode() return log_contents + except RepositoryError as err: + raise err except Exception as err: raise RetrievalError(f"Error retrieving the agent logs: {err}") diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py index d49d0957d..ac841096b 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -1,9 +1,16 @@ +import io +from unittest.mock import MagicMock from uuid import UUID import pytest from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from monkey_island.cc.repository import FileAgentLogRepository, RetrievalError +from monkey_island.cc.repository import ( + FileAgentLogRepository, + IFileRepository, + RetrievalError, + UnknownRecordError, +) LOG_CONTENTS = "lots of useful information" AGENT_ID = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") @@ -22,7 +29,7 @@ def test_store_agent_log(repository): def test_get_agent_log__unknown_record_error(repository): - with pytest.raises(RetrievalError): + with pytest.raises(UnknownRecordError): repository.get_agent_log(AGENT_ID) @@ -32,8 +39,17 @@ def test_get_agent_log__retrieval_error(): repository.get_agent_log(AGENT_ID) +def test_get_agent_log__corrupt_data(): + file_repository = MagicMock(spec=IFileRepository) + file_repository.open_file = MagicMock(return_value=io.BytesIO(b"\xff\xfe")) + repository = FileAgentLogRepository(file_repository) + + with pytest.raises(RetrievalError): + repository.get_agent_log(AGENT_ID) + + def test_reset_agent_logs(repository): repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) repository.reset() - with pytest.raises(RetrievalError): + with pytest.raises(UnknownRecordError): repository.get_agent_log(AGENT_ID) From 35d0cbc3b058453d3b8d19f22e041dd314b16682 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 12:55:58 -0400 Subject: [PATCH 23/28] UT: Add type hints to test_file_agent_log_repository.py --- .../cc/repository/test_file_agent_log_repository.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py index ac841096b..4ff7f904e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -7,6 +7,7 @@ from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository from monkey_island.cc.repository import ( FileAgentLogRepository, + IAgentLogRepository, IFileRepository, RetrievalError, UnknownRecordError, @@ -17,18 +18,18 @@ AGENT_ID = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") @pytest.fixture -def repository(): +def repository() -> IAgentLogRepository: return FileAgentLogRepository(SingleFileRepository()) -def test_store_agent_log(repository): +def test_store_agent_log(repository: IAgentLogRepository): repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) retrieved_log_contents = repository.get_agent_log(AGENT_ID) assert retrieved_log_contents == LOG_CONTENTS -def test_get_agent_log__unknown_record_error(repository): +def test_get_agent_log__unknown_record_error(repository: IAgentLogRepository): with pytest.raises(UnknownRecordError): repository.get_agent_log(AGENT_ID) @@ -48,7 +49,7 @@ def test_get_agent_log__corrupt_data(): repository.get_agent_log(AGENT_ID) -def test_reset_agent_logs(repository): +def test_reset_agent_logs(repository: IAgentLogRepository): repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) repository.reset() with pytest.raises(UnknownRecordError): From 0e2d82a7adfe67b77a13eab2f3356f550957bf49 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 12:57:21 -0400 Subject: [PATCH 24/28] UT: Add comment --- .../cc/repository/test_file_agent_log_repository.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py index 4ff7f904e..c36b43545 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -42,6 +42,7 @@ def test_get_agent_log__retrieval_error(): def test_get_agent_log__corrupt_data(): file_repository = MagicMock(spec=IFileRepository) + # Return invalid unicode file_repository.open_file = MagicMock(return_value=io.BytesIO(b"\xff\xfe")) repository = FileAgentLogRepository(file_repository) From 20d5fb3748dad5c91230c7529f37db5e1259f59a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 13:05:22 -0400 Subject: [PATCH 25/28] 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): From bddee026fe76849c45c835cfc006b6ff2086de77 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 13:29:34 -0400 Subject: [PATCH 26/28] UT: Add InMemoryFileRepository --- monkey/tests/monkey_island/__init__.py | 1 + .../in_memory_file_repository.py | 33 +++++++++++++++++++ 2 files changed, 34 insertions(+) create mode 100644 monkey/tests/monkey_island/in_memory_file_repository.py diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index f7033842d..9dd1996f9 100644 --- a/monkey/tests/monkey_island/__init__.py +++ b/monkey/tests/monkey_island/__init__.py @@ -4,3 +4,4 @@ from .open_error_file_repository import OpenErrorFileRepository from .in_memory_agent_configuration_repository import InMemoryAgentConfigurationRepository from .in_memory_simulation_configuration import InMemorySimulationRepository from .in_memory_credentials_repository import InMemoryCredentialsRepository +from .in_memory_file_repository import InMemoryFileRepository diff --git a/monkey/tests/monkey_island/in_memory_file_repository.py b/monkey/tests/monkey_island/in_memory_file_repository.py new file mode 100644 index 000000000..76778d226 --- /dev/null +++ b/monkey/tests/monkey_island/in_memory_file_repository.py @@ -0,0 +1,33 @@ +import io +import re +from typing import BinaryIO, Dict + +from common.utils.code_utils import del_key +from monkey_island.cc.repository import IFileRepository, UnknownRecordError + + +class InMemoryFileRepository(IFileRepository): + def __init__(self): + self._files: Dict[str, bytes] = {} + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + self._files[unsafe_file_name] = file_contents.read() + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + try: + return io.BytesIO(self._files[unsafe_file_name]) + except KeyError: + raise UnknownRecordError(f"Unknown file {unsafe_file_name}") + + def delete_file(self, unsafe_file_name: str): + del_key(self._files, "unsafe_file_name") + + def delete_files_by_regex(self, file_name_regex: re.Pattern): + self._files = { + name: contents + for name, contents in self._files.items() + if not re.match(file_name_regex, name) + } + + def delete_all_files(self): + self._files = {} From 412a58f1f2cb4b0bdb010f6b5a8ffebe903bf7a1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 13:30:01 -0400 Subject: [PATCH 27/28] UT: Use InMemoryFileRepository in test_file_agent_log_repository.py --- .../cc/repository/test_file_agent_log_repository.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py index c36b43545..6409e9eff 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock from uuid import UUID import pytest -from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository +from tests.monkey_island import InMemoryFileRepository, OpenErrorFileRepository from monkey_island.cc.repository import ( FileAgentLogRepository, @@ -19,7 +19,7 @@ AGENT_ID = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") @pytest.fixture def repository() -> IAgentLogRepository: - return FileAgentLogRepository(SingleFileRepository()) + return FileAgentLogRepository(InMemoryFileRepository()) def test_store_agent_log(repository: IAgentLogRepository): From 2e8afe218e70c151b0bee7eeff17781608c6d72d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 14:00:16 -0400 Subject: [PATCH 28/28] UT: Add test of multiple agent logs --- .../test_file_agent_log_repository.py | 28 +++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py index 6409e9eff..71a0f878e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_log_repository.py @@ -14,7 +14,8 @@ from monkey_island.cc.repository import ( ) LOG_CONTENTS = "lots of useful information" -AGENT_ID = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") +AGENT_ID_1 = UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8") +AGENT_ID_2 = UUID("789abcd4-20d7-abcd-ef7a-0123acaabcde") @pytest.fixture @@ -23,21 +24,21 @@ def repository() -> IAgentLogRepository: def test_store_agent_log(repository: IAgentLogRepository): - repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) - retrieved_log_contents = repository.get_agent_log(AGENT_ID) + repository.upsert_agent_log(AGENT_ID_1, LOG_CONTENTS) + retrieved_log_contents = repository.get_agent_log(AGENT_ID_1) assert retrieved_log_contents == LOG_CONTENTS def test_get_agent_log__unknown_record_error(repository: IAgentLogRepository): with pytest.raises(UnknownRecordError): - repository.get_agent_log(AGENT_ID) + repository.get_agent_log(AGENT_ID_1) def test_get_agent_log__retrieval_error(): repository = FileAgentLogRepository(OpenErrorFileRepository()) with pytest.raises(RetrievalError): - repository.get_agent_log(AGENT_ID) + repository.get_agent_log(AGENT_ID_1) def test_get_agent_log__corrupt_data(): @@ -47,11 +48,22 @@ def test_get_agent_log__corrupt_data(): repository = FileAgentLogRepository(file_repository) with pytest.raises(RetrievalError): - repository.get_agent_log(AGENT_ID) + repository.get_agent_log(AGENT_ID_1) + + +def test_multiple_logs(repository: IAgentLogRepository): + log_contents_1 = "hello" + log_contents_2 = "world" + + repository.upsert_agent_log(AGENT_ID_1, log_contents_1) + repository.upsert_agent_log(AGENT_ID_2, log_contents_2) + + assert repository.get_agent_log(AGENT_ID_1) == log_contents_1 + assert repository.get_agent_log(AGENT_ID_2) == log_contents_2 def test_reset_agent_logs(repository: IAgentLogRepository): - repository.upsert_agent_log(AGENT_ID, LOG_CONTENTS) + repository.upsert_agent_log(AGENT_ID_1, LOG_CONTENTS) repository.reset() with pytest.raises(UnknownRecordError): - repository.get_agent_log(AGENT_ID) + repository.get_agent_log(AGENT_ID_1)