Island: AgentLogs.GET to return 404 if not agent log is found

This commit is contained in:
Ilija Lazoroski 2022-09-28 16:30:28 +02:00
parent 96662f3f66
commit e4d45b25cb
2 changed files with 15 additions and 5 deletions

View File

@ -1,12 +1,15 @@
import logging
from http import HTTPStatus from http import HTTPStatus
from flask import request from flask import request
from common.types import AgentID 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.AbstractResource import AbstractResource
from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.resources.request_authentication import jwt_required
logger = logging.getLogger(__name__)
class AgentLogs(AbstractResource): class AgentLogs(AbstractResource):
urls = ["/api/agent-logs/<uuid:agent_id>"] urls = ["/api/agent-logs/<uuid:agent_id>"]
@ -16,13 +19,17 @@ class AgentLogs(AbstractResource):
@jwt_required @jwt_required
def get(self, agent_id: AgentID): 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 return log_contents, HTTPStatus.OK
def put(self, agent_id: AgentID): def put(self, agent_id: AgentID):
log_contents = request.json 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 return {}, HTTPStatus.NO_CONTENT

View File

@ -5,7 +5,7 @@ import pytest
from tests.common import StubDIContainer from tests.common import StubDIContainer
from common.types import AgentID 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_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5")
AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5")
@ -23,6 +23,8 @@ class StubAgentLogRepository(IAgentLogRepository):
self._agent_logs[agent_id] = log_contents self._agent_logs[agent_id] = log_contents
def get_agent_log(self, agent_id: AgentID) -> str: 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] return self._agent_logs[agent_id]
def reset(self): def reset(self):
@ -40,7 +42,8 @@ def flask_client(build_flask_client):
def test_agent_logs_endpoint__get_empty(flask_client): def test_agent_logs_endpoint__get_empty(flask_client):
resp = flask_client.get(AGENT_LOGS_URL_1, follow_redirects=True) 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( @pytest.mark.parametrize(