forked from p34709852/monkey
Island: Add `GET /api/agents/`
This commit is contained in:
parent
3b6e4f5313
commit
68b288e5b3
|
@ -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
|
||||
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue