BB: Use threading to download logs
Reduces time to download logs by approx. 40%, but may be unnecessary after resolving https://github.com/guardicore/monkey/issues/2383
This commit is contained in:
parent
e4155648c1
commit
6a783d9c3e
|
@ -1,6 +1,7 @@
|
|||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Mapping, Sequence
|
||||
from threading import Thread
|
||||
from typing import List, Mapping
|
||||
|
||||
from common.types import MachineID
|
||||
from monkey_island.cc.models import Agent, Machine
|
||||
|
@ -12,7 +13,7 @@ class MonkeyLogsDownloader(object):
|
|||
def __init__(self, island_client, log_dir_path):
|
||||
self.island_client = island_client
|
||||
self.log_dir_path = Path(log_dir_path)
|
||||
self.monkey_log_paths: Sequence[Path] = []
|
||||
self.monkey_log_paths: List[Path] = []
|
||||
|
||||
def download_monkey_logs(self):
|
||||
try:
|
||||
|
@ -21,16 +22,29 @@ class MonkeyLogsDownloader(object):
|
|||
agents = self.island_client.get_agents()
|
||||
machines = self.island_client.get_machines()
|
||||
|
||||
download_threads: List[Thread] = []
|
||||
|
||||
# TODO: Does downloading logs concurrently still improve performance after resolving
|
||||
# https://github.com/guardicore/monkey/issues/2383?
|
||||
for agent in agents:
|
||||
log_file_path = self._get_log_file_path(agent, machines)
|
||||
log_contents = self.island_client.get_agent_log(agent.id)
|
||||
t = Thread(target=self._download_log, args=(agent, machines), daemon=True)
|
||||
t.start()
|
||||
download_threads.append(t)
|
||||
|
||||
MonkeyLogsDownloader._write_log_to_file(log_file_path, log_contents)
|
||||
for thread in download_threads:
|
||||
thread.join()
|
||||
|
||||
self.monkey_log_paths.append(log_file_path)
|
||||
except Exception as err:
|
||||
LOGGER.exception(err)
|
||||
|
||||
def _download_log(self, agent: Agent, machines: Mapping[MachineID, Machine]):
|
||||
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)
|
||||
|
||||
def _get_log_file_path(self, agent: Agent, machines: Mapping[MachineID, Machine]) -> Path:
|
||||
try:
|
||||
machine_ip = machines[agent.machine_id].network_interfaces[0].ip
|
||||
|
|
Loading…
Reference in New Issue