forked from p15670423/monkey
Island: Add FileAgentLogRepository
This commit is contained in:
parent
ce0affb1ed
commit
1c6cfa1ce6
|
@ -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}*")
|
Loading…
Reference in New Issue