Island: Add `GET /api/agents/`

This commit is contained in:
Mike Salvatore 2022-09-30 11:05:26 -04:00
parent 3b6e4f5313
commit 68b288e5b3
3 changed files with 59 additions and 7 deletions

View File

@ -23,7 +23,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
- `/api/agent-events` endpoint. #2155, #2300
- The ability to customize the file extension used by ransomware when
encrypting files. #1242
- `/api/agents` endpoint.
- `/api/agents` endpoint. #2362
- `/api/agent-signals` endpoint. #2261
- `/api/agent-logs/<uuid:agent_id>` endpoint. #2274

View File

@ -6,6 +6,7 @@ from flask import make_response, request
from common import AgentRegistrationData
from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
from monkey_island.cc.repository import IAgentRepository
from monkey_island.cc.resources.AbstractResource import AbstractResource
logger = logging.getLogger(__name__)
@ -14,8 +15,12 @@ logger = logging.getLogger(__name__)
class Agents(AbstractResource):
urls = ["/api/agents"]
def __init__(self, island_event_queue: IIslandEventQueue):
def __init__(self, island_event_queue: IIslandEventQueue, agent_repository: IAgentRepository):
self._island_event_queue = island_event_queue
self._agent_repository = agent_repository
def get(self):
return self._agent_repository.get_agents(), HTTPStatus.OK
def post(self):
try:

View File

@ -6,11 +6,12 @@ import pytest
from tests.common import StubDIContainer
from tests.unit_tests.monkey_island.conftest import get_url_for_resource
from common.types import SocketAddress
from monkey_island.cc.event_queue import IIslandEventQueue
from monkey_island.cc.models import Agent
from monkey_island.cc.repository import IAgentRepository
from monkey_island.cc.resources import Agents
AGENTS_URL = get_url_for_resource(Agents)
AGENT_REGISTRATION_DICT = {
"id": UUID("6bfd8b64-43d8-4449-8c70-d898aca74ad8"),
"machine_hardware_id": 1,
@ -20,11 +21,37 @@ AGENT_REGISTRATION_DICT = {
"network_interfaces": ["10.1.1.2/24"],
}
AGENTS = (
Agent(
id=UUID("12345678-1234-1234-1234-123456789abc"),
machine_id=2,
start_time=0,
stop_time=10,
cc_server=SocketAddress(ip="10.0.0.1", port=5000),
),
Agent(
id=UUID("abcdef78-abcd-abcd-abcd-abcdef123456"),
machine_id=3,
start_time=5,
stop_time=15,
cc_server=SocketAddress(ip="10.0.0.1", port=5000),
),
)
@pytest.fixture
def flask_client(build_flask_client):
def agent_repository() -> IAgentRepository:
agent_repository = MagicMock(spec=IAgentRepository)
agent_repository.get_agents = MagicMock(return_value=AGENTS)
return agent_repository
@pytest.fixture
def flask_client(build_flask_client, agent_repository):
container = StubDIContainer()
container.register_instance(IIslandEventQueue, MagicMock(spec=IIslandEventQueue))
container.register_instance(IAgentRepository, agent_repository)
with build_flask_client(container) as flask_client:
yield flask_client
@ -32,7 +59,7 @@ def flask_client(build_flask_client):
def test_agent_registration(flask_client):
resp = flask_client.post(
AGENTS_URL,
get_url_for_resource(Agents),
json=AGENT_REGISTRATION_DICT,
follow_redirects=True,
)
@ -46,9 +73,29 @@ def test_agent_registration_invalid_data(flask_client):
agent_registration_dict["id"] = 1
resp = flask_client.post(
AGENTS_URL,
get_url_for_resource(Agents),
json=agent_registration_dict,
follow_redirects=True,
)
assert resp.status_code == HTTPStatus.BAD_REQUEST
def test_get_agents__status_code(flask_client):
resp = flask_client.get(
get_url_for_resource(Agents),
follow_redirects=True,
)
assert resp.status_code == HTTPStatus.OK
def test_get_agents__data(flask_client):
resp = flask_client.get(
get_url_for_resource(Agents),
follow_redirects=True,
)
agents = [Agent(**a) for a in resp.json]
assert len(agents) == len(AGENTS)
for a in agents:
assert a in AGENTS