From e4155648c108132a5c036666e5783ce36faa8248 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Sun, 2 Oct 2022 13:09:11 -0400 Subject: [PATCH] BB: Download agent logs from new endpoints --- .../log_handlers/monkey_logs_downloader.py | 50 ++++++++++++++----- 1 file changed, 37 insertions(+), 13 deletions(-) diff --git a/envs/monkey_zoo/blackbox/log_handlers/monkey_logs_downloader.py b/envs/monkey_zoo/blackbox/log_handlers/monkey_logs_downloader.py index 302da8fc7..96f0980ea 100644 --- a/envs/monkey_zoo/blackbox/log_handlers/monkey_logs_downloader.py +++ b/envs/monkey_zoo/blackbox/log_handlers/monkey_logs_downloader.py @@ -1,6 +1,9 @@ import logging +from pathlib import Path +from typing import Mapping, Sequence -from envs.monkey_zoo.blackbox.log_handlers.monkey_log import MonkeyLog +from common.types import MachineID +from monkey_island.cc.models import Agent, Machine LOGGER = logging.getLogger(__name__) @@ -8,18 +11,39 @@ LOGGER = logging.getLogger(__name__) class MonkeyLogsDownloader(object): def __init__(self, island_client, log_dir_path): self.island_client = island_client - self.log_dir_path = log_dir_path - self.monkey_log_paths = [] + self.log_dir_path = Path(log_dir_path) + self.monkey_log_paths: Sequence[Path] = [] def download_monkey_logs(self): - LOGGER.info("Downloading each monkey log.") - all_monkeys = self.island_client.get_all_monkeys_from_db() - for monkey in all_monkeys: - downloaded_log_path = self._download_monkey_log(monkey) - if downloaded_log_path: - self.monkey_log_paths.append(downloaded_log_path) + try: + LOGGER.info("Downloading each monkey log.") - def _download_monkey_log(self, monkey): - log_handler = MonkeyLog(monkey, self.log_dir_path) - download_successful = log_handler.download_log(self.island_client) - return log_handler.get_log_path_for_monkey(monkey) if download_successful else None + agents = self.island_client.get_agents() + machines = self.island_client.get_machines() + + for agent in agents: + log_file_path = self._get_log_file_path(agent, machines) + log_contents = self.island_client.get_agent_log(agent.id) + + MonkeyLogsDownloader._write_log_to_file(log_file_path, log_contents) + + self.monkey_log_paths.append(log_file_path) + except Exception as err: + LOGGER.exception(err) + + def _get_log_file_path(self, agent: Agent, machines: Mapping[MachineID, Machine]) -> Path: + try: + machine_ip = machines[agent.machine_id].network_interfaces[0].ip + except IndexError: + machine_ip = "UNKNOWN" + + start_time = agent.start_time.strftime("%Y-%m-%d-%H-%M-%S") + + return self.log_dir_path / f"agent-{start_time}-{machine_ip}.log" + + @staticmethod + def _write_log_to_file(log_file_path: Path, log_contents: str): + LOGGER.debug(f"Writing {len(log_contents)} bytes to {log_file_path}") + + with open(log_file_path, "w") as f: + f.write(log_contents)