forked from p15670423/monkey
Island: AgentLogs.GET to return 404 if not agent log is found
This commit is contained in:
parent
96662f3f66
commit
e4d45b25cb
|
@ -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):
|
||||||
|
try:
|
||||||
log_contents = self._agent_log_repository.get_agent_log(agent_id)
|
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
|
||||||
|
|
|
@ -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(
|
||||||
|
|
Loading…
Reference in New Issue