From 99c2c5c6ef4c66e844ba1c40de996471ea3c9196 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
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)