forked from p15670423/monkey
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,
|
RetrievalError,
|
||||||
)
|
)
|
||||||
|
|
||||||
AGENT_LOG_FILE_NAME_PATTERN = "agent-*.log"
|
|
||||||
AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$")
|
AGENT_LOG_FILE_NAME_REGEX = re.compile(r"^agent-[\w-]+.log$")
|
||||||
|
|
||||||
|
|
||||||
|
@ -18,13 +17,14 @@ class FileAgentLogRepository(IAgentLogRepository):
|
||||||
self._file_repository = file_repository
|
self._file_repository = file_repository
|
||||||
|
|
||||||
def upsert_agent_log(self, agent_id: AgentID, log_contents: str):
|
def upsert_agent_log(self, agent_id: AgentID, log_contents: str):
|
||||||
self._file_repository.save_file(
|
log_file_name = FileAgentLogRepository._get_agent_log_file_name(agent_id)
|
||||||
self._get_agent_log_file_name(agent_id), io.BytesIO(log_contents.encode())
|
self._file_repository.save_file(log_file_name, io.BytesIO(log_contents.encode()))
|
||||||
)
|
|
||||||
|
|
||||||
def get_agent_log(self, agent_id: AgentID) -> str:
|
def get_agent_log(self, agent_id: AgentID) -> str:
|
||||||
|
log_file_name = FileAgentLogRepository._get_agent_log_file_name(agent_id)
|
||||||
|
|
||||||
try:
|
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()
|
log_contents = f.read().decode()
|
||||||
return log_contents
|
return log_contents
|
||||||
except RepositoryError as err:
|
except RepositoryError as err:
|
||||||
|
@ -35,5 +35,6 @@ class FileAgentLogRepository(IAgentLogRepository):
|
||||||
def reset(self):
|
def reset(self):
|
||||||
self._file_repository.delete_files_by_regex(AGENT_LOG_FILE_NAME_REGEX)
|
self._file_repository.delete_files_by_regex(AGENT_LOG_FILE_NAME_REGEX)
|
||||||
|
|
||||||
def _get_agent_log_file_name(self, agent_id: AgentID) -> str:
|
@staticmethod
|
||||||
return AGENT_LOG_FILE_NAME_PATTERN.replace("*", str(agent_id))
|
def _get_agent_log_file_name(agent_id: AgentID) -> str:
|
||||||
|
return f"agent-{agent_id}.log"
|
||||||
|
|
|
@ -1,4 +1,5 @@
|
||||||
import io
|
import io
|
||||||
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
@ -36,7 +37,7 @@ class FileRepositoryCachingDecorator(IFileRepository):
|
||||||
self._open_file.cache_clear()
|
self._open_file.cache_clear()
|
||||||
return self._file_repository.delete_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):
|
||||||
self._open_file.cache_clear()
|
self._open_file.cache_clear()
|
||||||
return self._file_repository.delete_files_by_regex(file_name_regex)
|
return self._file_repository.delete_files_by_regex(file_name_regex)
|
||||||
|
|
||||||
|
|
|
@ -1,3 +1,4 @@
|
||||||
|
import re
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
from readerwriterlock import rwlock
|
from readerwriterlock import rwlock
|
||||||
|
@ -26,7 +27,7 @@ class FileRepositoryLockingDecorator(IFileRepository):
|
||||||
with self._rwlock.gen_wlock():
|
with self._rwlock.gen_wlock():
|
||||||
return self._file_repository.delete_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):
|
||||||
with self._rwlock.gen_wlock():
|
with self._rwlock.gen_wlock():
|
||||||
return self._file_repository.delete_files_by_regex(file_name_regex)
|
return self._file_repository.delete_files_by_regex(file_name_regex)
|
||||||
|
|
||||||
|
|
|
@ -1,4 +1,5 @@
|
||||||
import logging
|
import logging
|
||||||
|
import re
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
from . import IFileRepository
|
from . import IFileRepository
|
||||||
|
@ -26,7 +27,7 @@ class FileRepositoryLoggingDecorator(IFileRepository):
|
||||||
logger.debug(f"Deleting file {unsafe_file_name}")
|
logger.debug(f"Deleting file {unsafe_file_name}")
|
||||||
return self._file_repository.delete_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}"')
|
logger.debug(f'Deleting files whose names match the regex "{file_name_regex}"')
|
||||||
return self._file_repository.delete_files_by_regex(file_name_regex)
|
return self._file_repository.delete_files_by_regex(file_name_regex)
|
||||||
|
|
||||||
|
|
|
@ -1,4 +1,5 @@
|
||||||
import abc
|
import abc
|
||||||
|
import re
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
from monkey_island.cc.repository import UnknownRecordError
|
from monkey_island.cc.repository import UnknownRecordError
|
||||||
|
@ -51,7 +52,7 @@ class IFileRepository(metaclass=abc.ABCMeta):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@abc.abstractmethod
|
@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
|
Delete files whose names match a particular regex
|
||||||
|
|
||||||
|
|
|
@ -1,4 +1,5 @@
|
||||||
import io
|
import io
|
||||||
|
import re
|
||||||
from typing import BinaryIO
|
from typing import BinaryIO
|
||||||
|
|
||||||
from monkey_island.cc import repository
|
from monkey_island.cc import repository
|
||||||
|
@ -24,7 +25,7 @@ class MockFileRepository(IFileRepository):
|
||||||
def delete_file(self, unsafe_file_name: str):
|
def delete_file(self, unsafe_file_name: str):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def delete_files_by_regex(self, file_name_regex: str):
|
def delete_files_by_regex(self, file_name_regex: re.Pattern):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def delete_all_files(self):
|
def delete_all_files(self):
|
||||||
|
|
Loading…
Reference in New Issue