BB: Download agent logs from new endpoints

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.log_dir_path = Path(log_dir_path)
-        self.monkey_log_paths = []
+        self.monkey_log_paths: Sequence[Path] = []
 
     def download_monkey_logs(self):
-        LOGGER.info("Downloading each monkey log.")
+        try:
-        all_monkeys = self.island_client.get_all_monkeys_from_db()
+            LOGGER.info("Downloading each monkey log.")
-        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)
 
-    def _download_monkey_log(self, monkey):
+            agents = self.island_client.get_agents()
-        log_handler = MonkeyLog(monkey, self.log_dir_path)
+            machines = self.island_client.get_machines()
-        download_successful = log_handler.download_log(self.island_client)
+
-        return log_handler.get_log_path_for_monkey(monkey) if download_successful else None
+            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)