Island: Use re.Pattern for IFileRepository.delete_files_by_regex()
This commit is contained in:
parent
0e2d82a7ad
commit
20d5fb3748
|
@ -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"
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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):
|
||||
|
|
Loading…
Reference in New Issue