From 99c2c5c6ef4c66e844ba1c40de996471ea3c9196 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Sat, 1 Oct 2022 19:51:13 -0400 Subject: [PATCH] BB: Add MonkeyIslandClient.get_machines() --- .../blackbox/island_client/monkey_island_client.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py index 78c2441be..85405af2d 100644 --- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py +++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py @@ -1,17 +1,19 @@ import json import logging import time -from typing import List, Sequence, Union +from typing import List, Mapping, Sequence, Union from bson import json_util from common.credentials import Credentials +from common.types import MachineID from envs.monkey_zoo.blackbox.island_client.monkey_island_requests import MonkeyIslandRequests from envs.monkey_zoo.blackbox.test_configurations.test_configuration import TestConfiguration -from monkey_island.cc.models import Agent +from monkey_island.cc.models import Agent, Machine SLEEP_BETWEEN_REQUESTS_SECONDS = 0.5 GET_AGENTS_ENDPOINT = "api/agents" +GET_MACHINES_ENDPOINT = "api/machines" MONKEY_TEST_ENDPOINT = "api/test/monkey" TELEMETRY_TEST_ENDPOINT = "api/test/telemetry" LOG_TEST_ENDPOINT = "api/test/log" @@ -164,6 +166,12 @@ class MonkeyIslandClient(object): return [Agent(**a) for a in response.json()] + def get_machines(self) -> Mapping[MachineID, Machine]: + response = self.requests.get(GET_MACHINES_ENDPOINT) + machines = (Machine(**m) for m in response.json()) + + return {m.id: m for m in machines} + def find_log_in_db(self, query): response = self.requests.get( LOG_TEST_ENDPOINT, MonkeyIslandClient.form_find_query_for_request(query)