From 543c063f7bb48ffff6bdac20f1b19af495604218 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 14:52:19 +0200 Subject: [PATCH 01/14] Island: Add AgentLogs endpoint --- monkey/monkey_island/cc/resources/__init__.py | 1 + .../monkey_island/cc/resources/agent_logs.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 monkey/monkey_island/cc/resources/agent_logs.py diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index 937766e2b..e6c2db648 100644 --- a/monkey/monkey_island/cc/resources/__init__.py +++ b/monkey/monkey_island/cc/resources/__init__.py @@ -11,3 +11,4 @@ from .pba_file_download import PBAFileDownload from .agent_events import AgentEvents from .agents import Agents from .agent_signals import AgentSignals, TerminateAllAgents +from .agent_logs import AgentLogs diff --git a/monkey/monkey_island/cc/resources/agent_logs.py b/monkey/monkey_island/cc/resources/agent_logs.py new file mode 100644 index 000000000..64270fd64 --- /dev/null +++ b/monkey/monkey_island/cc/resources/agent_logs.py @@ -0,0 +1,27 @@ +from http import HTTPStatus + +from flask import request + +from monkey_island.cc.repository import IAgentLogRepository +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required + + +class AgentLogs(AbstractResource): + urls = ["/api/agent-logs/"] + + def __init__(self, agent_log_repository: IAgentLogRepository): + self._agent_log_repository = agent_log_repository + + @jwt_required + def get(self, agent_id: str): + agent_log = self._agent_log_repository.get_agent_log(agent_id) + + return agent_log, HTTPStatus.OK + + def put(self, agent_id: str): + agent_data = request.json["log_contents"] + + self._agent_log_repository.upsert_agent_log(agent_id, agent_data) + + return {}, HTTPStatus.NO_CONTENT From 89ae9824d0614850790ebf79718c3c9f01a474ea Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 14:53:24 +0200 Subject: [PATCH 02/14] Island: Add the new AgentLogs endpoint to the island --- monkey/monkey_island/cc/app.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 494f48dd0..5e0aaadb2 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -14,6 +14,7 @@ from monkey_island.cc.resources import ( AgentBinaries, AgentConfiguration, AgentEvents, + AgentLogs, Agents, AgentSignals, ClearSimulationData, @@ -185,6 +186,7 @@ def init_restful_endpoints(api: FlaskDIWrapper): api.add_resource(ZeroTrustFindingEvent) api.add_resource(TelemetryFeed) api.add_resource(Log) + api.add_resource(AgentLogs) api.add_resource(IslandLog) api.add_resource(IPAddresses) From 546c44f501f56c9b3e3f4dc421995cd224c28e8d Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 14:55:37 +0200 Subject: [PATCH 03/14] UT: Add tests for AgentLogs endpoint --- .../cc/resources/test_agent_logs.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py new file mode 100644 index 000000000..ead1c73a2 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py @@ -0,0 +1,59 @@ +from http import HTTPStatus +from uuid import UUID + +import pytest +from tests.common import StubDIContainer + +from common.types import AgentID +from monkey_island.cc.repository import IAgentLogRepository + +AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") +AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") + +AGENT_LOGS_URL_1 = f"/api/agent-logs/{AGENT_ID_1}" +AGENT_LOGS_URL_2 = f"/api/agent-logs/{AGENT_ID_2}" + + +class StubAgentLogRepository(IAgentLogRepository): + def __init__(self): + self._agent_logs = {} + + def upsert_agent_log(self, agent_id: AgentID, log_contents: str): + if agent_id not in self._agent_logs.keys(): + self._agent_logs[agent_id] = log_contents + + def get_agent_log(self, agent_id: AgentID) -> str: + return self._agent_logs[agent_id] + + def reset(self): + self._agent_logs = {} + + +@pytest.fixture +def flask_client(build_flask_client): + container = StubDIContainer() + container.register_instance(IAgentLogRepository, StubAgentLogRepository()) + + with build_flask_client(container) as flask_client: + yield flask_client + + +def test_agent_logs_endpoint__get_empty(flask_client): + resp = flask_client.get(AGENT_LOGS_URL_1, follow_redirects=True) + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + + +@pytest.mark.parametrize( + "url,log", [(AGENT_LOGS_URL_1, "LoremIpsum1"), (AGENT_LOGS_URL_2, "SecondLoremIpsum")] +) +def test_agent_logs_endpoint(flask_client, url, log): + flask_client.put(url, json={"log_contents": log}, follow_redirects=True) + resp = flask_client.get(url, follow_redirects=True) + assert resp.status_code == HTTPStatus.OK + assert resp.json == log + + +def test_agent_logs_endpoint__put_bad_json(flask_client): + resp = flask_client.put(AGENT_LOGS_URL_1, json={"bad_key": "bad"}, follow_redirects=True) + + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR From 08bc43e0c46ed94d9fb89032b2dc99162fc373be Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 14:57:09 +0200 Subject: [PATCH 04/14] Changelog: Add entry for adding new AgentLogs endpoint --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 72782b3e4..2e50b2959 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). encrypting files. #1242 - `/api/agents` endpoint. - `/api/agent-signals` endpoint. #2261 +- `/api/agent-logs/` endpoint. #2274 ### Changed - Reset workflow. Now it's possible to delete data gathered by agents without From fa13ca8df8f1954faf4eb164cc2c96290efd65a4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 09:31:38 -0400 Subject: [PATCH 05/14] Island: Rename agent_log -> log_contents --- monkey/monkey_island/cc/resources/agent_logs.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_logs.py b/monkey/monkey_island/cc/resources/agent_logs.py index 64270fd64..b9c7b1957 100644 --- a/monkey/monkey_island/cc/resources/agent_logs.py +++ b/monkey/monkey_island/cc/resources/agent_logs.py @@ -15,13 +15,12 @@ class AgentLogs(AbstractResource): @jwt_required def get(self, agent_id: str): - agent_log = self._agent_log_repository.get_agent_log(agent_id) + log_contents = self._agent_log_repository.get_agent_log(agent_id) - return agent_log, HTTPStatus.OK + return log_contents, HTTPStatus.OK def put(self, agent_id: str): - agent_data = request.json["log_contents"] - - self._agent_log_repository.upsert_agent_log(agent_id, agent_data) + log_contents = request.json["log_contents"] + self._agent_log_repository.upsert_agent_log(agent_id, log_contents) return {}, HTTPStatus.NO_CONTENT From b705e33af390897595fe6011cd55615115edf24e Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:21:39 +0200 Subject: [PATCH 06/14] Island: Add AgentID type for agent_id to AgentLogs endpoint --- monkey/monkey_island/cc/resources/agent_logs.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_logs.py b/monkey/monkey_island/cc/resources/agent_logs.py index b9c7b1957..e6bb287a9 100644 --- a/monkey/monkey_island/cc/resources/agent_logs.py +++ b/monkey/monkey_island/cc/resources/agent_logs.py @@ -2,25 +2,27 @@ from http import HTTPStatus from flask import request +from common.types import AgentID from monkey_island.cc.repository import IAgentLogRepository from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required class AgentLogs(AbstractResource): - urls = ["/api/agent-logs/"] + urls = ["/api/agent-logs/"] def __init__(self, agent_log_repository: IAgentLogRepository): self._agent_log_repository = agent_log_repository @jwt_required - def get(self, agent_id: str): + def get(self, agent_id: AgentID): log_contents = self._agent_log_repository.get_agent_log(agent_id) return log_contents, HTTPStatus.OK - def put(self, agent_id: str): - log_contents = request.json["log_contents"] - self._agent_log_repository.upsert_agent_log(agent_id, log_contents) + def put(self, agent_id: AgentID): + log_contents = request.json + + self._agent_log_repository.upsert_agent_log(agent_id, agent_data) return {}, HTTPStatus.NO_CONTENT From 96662f3f669442cb84ba144282cbc8b3860608e1 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:25:14 +0200 Subject: [PATCH 07/14] UT: Fix AgentLogs tests --- .../monkey_island/cc/resources/test_agent_logs.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py index ead1c73a2..42249cf46 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py @@ -47,13 +47,7 @@ def test_agent_logs_endpoint__get_empty(flask_client): "url,log", [(AGENT_LOGS_URL_1, "LoremIpsum1"), (AGENT_LOGS_URL_2, "SecondLoremIpsum")] ) def test_agent_logs_endpoint(flask_client, url, log): - flask_client.put(url, json={"log_contents": log}, follow_redirects=True) + flask_client.put(url, json=log, follow_redirects=True) resp = flask_client.get(url, follow_redirects=True) assert resp.status_code == HTTPStatus.OK assert resp.json == log - - -def test_agent_logs_endpoint__put_bad_json(flask_client): - resp = flask_client.put(AGENT_LOGS_URL_1, json={"bad_key": "bad"}, follow_redirects=True) - - assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR From e4d45b25cbebb0603f5098b08180910786a228d6 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:30:28 +0200 Subject: [PATCH 08/14] Island: AgentLogs.GET to return 404 if not agent log is found --- monkey/monkey_island/cc/resources/agent_logs.py | 13 ++++++++++--- .../monkey_island/cc/resources/test_agent_logs.py | 7 +++++-- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_logs.py b/monkey/monkey_island/cc/resources/agent_logs.py index e6bb287a9..16c4992fa 100644 --- a/monkey/monkey_island/cc/resources/agent_logs.py +++ b/monkey/monkey_island/cc/resources/agent_logs.py @@ -1,12 +1,15 @@ +import logging from http import HTTPStatus from flask import request from common.types import AgentID -from monkey_island.cc.repository import IAgentLogRepository +from monkey_island.cc.repository import IAgentLogRepository, UnknownRecordError from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required +logger = logging.getLogger(__name__) + class AgentLogs(AbstractResource): urls = ["/api/agent-logs/"] @@ -16,13 +19,17 @@ class AgentLogs(AbstractResource): @jwt_required def get(self, agent_id: AgentID): - log_contents = self._agent_log_repository.get_agent_log(agent_id) + try: + log_contents = self._agent_log_repository.get_agent_log(agent_id) + except UnknownRecordError as err: + logger.debug(f"Error occured while getting agent log: {err}") + return {}, HTTPStatus.NOT_FOUND return log_contents, HTTPStatus.OK def put(self, agent_id: AgentID): log_contents = request.json - self._agent_log_repository.upsert_agent_log(agent_id, agent_data) + self._agent_log_repository.upsert_agent_log(agent_id, log_contents) return {}, HTTPStatus.NO_CONTENT diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py index 42249cf46..f6eb81e65 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py @@ -5,7 +5,7 @@ import pytest from tests.common import StubDIContainer from common.types import AgentID -from monkey_island.cc.repository import IAgentLogRepository +from monkey_island.cc.repository import IAgentLogRepository, UnknownRecordError AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") @@ -23,6 +23,8 @@ class StubAgentLogRepository(IAgentLogRepository): self._agent_logs[agent_id] = log_contents def get_agent_log(self, agent_id: AgentID) -> str: + if agent_id not in self._agent_logs: + raise UnknownRecordError("Error occured while getting agent") return self._agent_logs[agent_id] def reset(self): @@ -40,7 +42,8 @@ def flask_client(build_flask_client): def test_agent_logs_endpoint__get_empty(flask_client): resp = flask_client.get(AGENT_LOGS_URL_1, follow_redirects=True) - assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert resp.status_code == HTTPStatus.NOT_FOUND + assert resp.json == {} @pytest.mark.parametrize( From d0293b4edc899ab29ae5aed76b90d217d347f6e9 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:36:38 +0200 Subject: [PATCH 09/14] UT: Move and rename StubAgentLogRepository to InMemoryAgentLogRepository --- monkey/tests/monkey_island/__init__.py | 1 + .../in_memory_agent_log_repository.py | 19 +++++++++++++++ .../cc/resources/test_agent_logs.py | 23 +++---------------- 3 files changed, 23 insertions(+), 20 deletions(-) create mode 100644 monkey/tests/monkey_island/in_memory_agent_log_repository.py diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index f7033842d..4b7f0f7c7 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_agent_log_repository import InMemoryAgentLogRepository diff --git a/monkey/tests/monkey_island/in_memory_agent_log_repository.py b/monkey/tests/monkey_island/in_memory_agent_log_repository.py new file mode 100644 index 000000000..ccd1d18cf --- /dev/null +++ b/monkey/tests/monkey_island/in_memory_agent_log_repository.py @@ -0,0 +1,19 @@ +from common.types import AgentID +from monkey_island.cc.repository import IAgentLogRepository, UnknownRecordError + + +class InMemoryAgentLogRepository(IAgentLogRepository): + def __init__(self): + self._agent_logs = {} + + def upsert_agent_log(self, agent_id: AgentID, log_contents: str): + if agent_id not in self._agent_logs.keys(): + self._agent_logs[agent_id] = log_contents + + def get_agent_log(self, agent_id: AgentID) -> str: + if agent_id not in self._agent_logs: + raise UnknownRecordError("Error occured while getting agent") + return self._agent_logs[agent_id] + + def reset(self): + self._agent_logs = {} diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py index f6eb81e65..0f7e86b53 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py @@ -3,9 +3,9 @@ from uuid import UUID import pytest from tests.common import StubDIContainer +from tests.monkey_island import InMemoryAgentLogRepository -from common.types import AgentID -from monkey_island.cc.repository import IAgentLogRepository, UnknownRecordError +from monkey_island.cc.repository import IAgentLogRepository AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") @@ -14,27 +14,10 @@ AGENT_LOGS_URL_1 = f"/api/agent-logs/{AGENT_ID_1}" AGENT_LOGS_URL_2 = f"/api/agent-logs/{AGENT_ID_2}" -class StubAgentLogRepository(IAgentLogRepository): - def __init__(self): - self._agent_logs = {} - - def upsert_agent_log(self, agent_id: AgentID, log_contents: str): - if agent_id not in self._agent_logs.keys(): - self._agent_logs[agent_id] = log_contents - - def get_agent_log(self, agent_id: AgentID) -> str: - if agent_id not in self._agent_logs: - raise UnknownRecordError("Error occured while getting agent") - return self._agent_logs[agent_id] - - def reset(self): - self._agent_logs = {} - - @pytest.fixture def flask_client(build_flask_client): container = StubDIContainer() - container.register_instance(IAgentLogRepository, StubAgentLogRepository()) + container.register_instance(IAgentLogRepository, InMemoryAgentLogRepository()) with build_flask_client(container) as flask_client: yield flask_client From 1a8306af1b043ac756efa9f077c20a0adfa72def Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:43:33 +0200 Subject: [PATCH 10/14] Island: Change log level in AgentLogs endpoint --- monkey/monkey_island/cc/resources/agent_logs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/agent_logs.py b/monkey/monkey_island/cc/resources/agent_logs.py index 16c4992fa..1426df792 100644 --- a/monkey/monkey_island/cc/resources/agent_logs.py +++ b/monkey/monkey_island/cc/resources/agent_logs.py @@ -22,7 +22,7 @@ class AgentLogs(AbstractResource): try: log_contents = self._agent_log_repository.get_agent_log(agent_id) except UnknownRecordError as err: - logger.debug(f"Error occured while getting agent log: {err}") + logger.error(f"Error occured while getting agent log: {err}") return {}, HTTPStatus.NOT_FOUND return log_contents, HTTPStatus.OK From 10954e0a6eb47a59c17b9e9ef580126c4adb7db6 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:45:28 +0200 Subject: [PATCH 11/14] Island: Use UUID for AgentSignals endpoint parameter --- .../cc/resources/agent_signals/agent_signals.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals/agent_signals.py index f6ce22f71..1103722b5 100644 --- a/monkey/monkey_island/cc/resources/agent_signals/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals/agent_signals.py @@ -1,6 +1,7 @@ import logging from http import HTTPStatus +from common.types import AgentID from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services import AgentSignalsService @@ -8,7 +9,7 @@ logger = logging.getLogger(__name__) class AgentSignals(AbstractResource): - urls = ["/api/agent-signals/"] + urls = ["/api/agent-signals/"] def __init__( self, @@ -16,6 +17,6 @@ class AgentSignals(AbstractResource): ): self._agent_signals_service = agent_signals_service - def get(self, agent_id: str): + def get(self, agent_id: AgentID): agent_signals = self._agent_signals_service.get_signals(agent_id) return agent_signals.dict(simplify=True), HTTPStatus.OK From c69a414a4d519c9679825400f6a1d73c1464b6be Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 14:06:09 -0400 Subject: [PATCH 12/14] Island: Fix typo occured -> occurred --- monkey/monkey_island/cc/resources/agent_logs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/agent_logs.py b/monkey/monkey_island/cc/resources/agent_logs.py index 1426df792..40152c0fb 100644 --- a/monkey/monkey_island/cc/resources/agent_logs.py +++ b/monkey/monkey_island/cc/resources/agent_logs.py @@ -22,7 +22,7 @@ class AgentLogs(AbstractResource): try: log_contents = self._agent_log_repository.get_agent_log(agent_id) except UnknownRecordError as err: - logger.error(f"Error occured while getting agent log: {err}") + logger.error(f"Error occurred while getting agent log: {err}") return {}, HTTPStatus.NOT_FOUND return log_contents, HTTPStatus.OK From 3202404e46369e743c923466a080ee3ee96ad043 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 14:06:50 -0400 Subject: [PATCH 13/14] Changelog: Update agent-logs to use uuid instead of str This was missed in 10954e0a6e. --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e50b2959..57c7dbe75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). encrypting files. #1242 - `/api/agents` endpoint. - `/api/agent-signals` endpoint. #2261 -- `/api/agent-logs/` endpoint. #2274 +- `/api/agent-logs/` endpoint. #2274 ### Changed - Reset workflow. Now it's possible to delete data gathered by agents without From 93b0fe0f6ebaf7365002b7410030402ce83171c4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 28 Sep 2022 14:10:30 -0400 Subject: [PATCH 14/14] UT: Fix UnknownRecordError message in InMemoryAgentLogRepository --- monkey/tests/monkey_island/in_memory_agent_log_repository.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/tests/monkey_island/in_memory_agent_log_repository.py b/monkey/tests/monkey_island/in_memory_agent_log_repository.py index ccd1d18cf..fb75ad7a7 100644 --- a/monkey/tests/monkey_island/in_memory_agent_log_repository.py +++ b/monkey/tests/monkey_island/in_memory_agent_log_repository.py @@ -12,7 +12,8 @@ class InMemoryAgentLogRepository(IAgentLogRepository): def get_agent_log(self, agent_id: AgentID) -> str: if agent_id not in self._agent_logs: - raise UnknownRecordError("Error occured while getting agent") + raise UnknownRecordError(f"Unknown agent {agent_id}") + return self._agent_logs[agent_id] def reset(self):