From e4d45b25cbebb0603f5098b08180910786a228d6 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 28 Sep 2022 16:30:28 +0200 Subject: [PATCH] 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(