BB: Add MonkeyIslandClient.get_machines()

This commit is contained in:
Mike Salvatore 2022-10-01 19:51:13 -04:00
parent b335601a05
commit 99c2c5c6ef
1 changed files with 10 additions and 2 deletions

View File

@ -1,17 +1,19 @@
import json import json
import logging import logging
import time import time
from typing import List, Sequence, Union from typing import List, Mapping, Sequence, Union
from bson import json_util from bson import json_util
from common.credentials import Credentials 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.island_client.monkey_island_requests import MonkeyIslandRequests
from envs.monkey_zoo.blackbox.test_configurations.test_configuration import TestConfiguration 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 SLEEP_BETWEEN_REQUESTS_SECONDS = 0.5
GET_AGENTS_ENDPOINT = "api/agents" GET_AGENTS_ENDPOINT = "api/agents"
GET_MACHINES_ENDPOINT = "api/machines"
MONKEY_TEST_ENDPOINT = "api/test/monkey" MONKEY_TEST_ENDPOINT = "api/test/monkey"
TELEMETRY_TEST_ENDPOINT = "api/test/telemetry" TELEMETRY_TEST_ENDPOINT = "api/test/telemetry"
LOG_TEST_ENDPOINT = "api/test/log" LOG_TEST_ENDPOINT = "api/test/log"
@ -164,6 +166,12 @@ class MonkeyIslandClient(object):
return [Agent(**a) for a in response.json()] 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): def find_log_in_db(self, query):
response = self.requests.get( response = self.requests.get(
LOG_TEST_ENDPOINT, MonkeyIslandClient.form_find_query_for_request(query) LOG_TEST_ENDPOINT, MonkeyIslandClient.form_find_query_for_request(query)