forked from p15670423/monkey
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
|
import logging
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Mapping, Sequence
|
from threading import Thread
|
||||||
|
from typing import List, Mapping
|
||||||
|
|
||||||
from common.types import MachineID
|
from common.types import MachineID
|
||||||
from monkey_island.cc.models import Agent, Machine
|
from monkey_island.cc.models import Agent, Machine
|
||||||
|
@ -12,7 +13,7 @@ class MonkeyLogsDownloader(object):
|
||||||
def __init__(self, island_client, log_dir_path):
|
def __init__(self, island_client, log_dir_path):
|
||||||
self.island_client = island_client
|
self.island_client = island_client
|
||||||
self.log_dir_path = Path(log_dir_path)
|
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):
|
def download_monkey_logs(self):
|
||||||
try:
|
try:
|
||||||
|
@ -21,16 +22,29 @@ class MonkeyLogsDownloader(object):
|
||||||
agents = self.island_client.get_agents()
|
agents = self.island_client.get_agents()
|
||||||
machines = self.island_client.get_machines()
|
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:
|
for agent in agents:
|
||||||
log_file_path = self._get_log_file_path(agent, machines)
|
t = Thread(target=self._download_log, args=(agent, machines), daemon=True)
|
||||||
log_contents = self.island_client.get_agent_log(agent.id)
|
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:
|
except Exception as err:
|
||||||
LOGGER.exception(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:
|
def _get_log_file_path(self, agent: Agent, machines: Mapping[MachineID, Machine]) -> Path:
|
||||||
try:
|
try:
|
||||||
machine_ip = machines[agent.machine_id].network_interfaces[0].ip
|
machine_ip = machines[agent.machine_id].network_interfaces[0].ip
|
||||||
|
|
Loading…
Reference in New Issue