From baa1598a1b64ba19d9cd5fde1a8f8c392464cd40 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Fri, 24 Apr 2020 13:19:07 +0300
Subject: [PATCH] Added progress bars, optimized imports, fixed bug that
 happened on big numbers of telems

---
 .gitignore                                         |  1 -
 envs/monkey_zoo/blackbox/README.md                 |  4 ++--
 .../blackbox/analyzers/performance_analyzer.py     |  1 +
 .../blackbox/island_client/monkey_island_client.py |  5 ++---
 .../island_client/monkey_island_requests.py        |  2 +-
 .../monkey_zoo/blackbox/log_handlers/monkey_log.py |  2 +-
 .../blackbox/log_handlers/test_logs_handler.py     |  3 +--
 envs/monkey_zoo/blackbox/test_blackbox.py          | 14 ++++++--------
 envs/monkey_zoo/blackbox/tests/exploitation.py     |  5 ++---
 .../tests/performance/endpoint_performance_test.py |  7 +++----
 .../blackbox/tests/performance/map_generation.py   |  2 +-
 .../tests/performance/performance_test_workflow.py |  2 +-
 .../tests/performance/report_generation.py         |  2 +-
 .../performance/telemetry_performance_test.py      | 12 +++++++-----
 .../tests/performance/utils/telem_parser.py        | 13 ++++++++-----
 .../blackbox/utils/gcp_machine_handlers.py         |  3 +--
 envs/monkey_zoo/blackbox/utils/json_encoder.py     |  1 +
 .../cc/resources/test/utils/telem_store.py         |  4 +++-
 monkey/monkey_island/requirements.txt              |  1 +
 19 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/.gitignore b/.gitignore
index 71aa824de..97e618d43 100644
--- a/.gitignore
+++ b/.gitignore
@@ -85,6 +85,5 @@ MonkeyZoo/*
 # Exported monkey telemetries
 /monkey/test_telems/
 
-
 # vim swap files
 *.swp
diff --git a/envs/monkey_zoo/blackbox/README.md b/envs/monkey_zoo/blackbox/README.md
index 278ca65ef..334dbf3ac 100644
--- a/envs/monkey_zoo/blackbox/README.md
+++ b/envs/monkey_zoo/blackbox/README.md
@@ -12,10 +12,10 @@ this information in the GCP Console `Compute Engine/VM Instances` under _Externa
 #### Running in command line
 Run the following command:
 
-`monkey\envs\monkey_zoo\blackbox>python -m pytest --island=35.207.152.72:5000 test_blackbox.py`
+`monkey\envs\monkey_zoo\blackbox>python -m pytest -s --island=35.207.152.72:5000 test_blackbox.py`
 
 #### Running in PyCharm
-Configure a PyTest configuration with the additional argument `--island=35.207.152.72` on the 
+Configure a PyTest configuration with the additional arguments `-s --island=35.207.152.72` on the 
 `monkey\envs\monkey_zoo\blackbox`.
 
 ### Running telemetry performance test
diff --git a/envs/monkey_zoo/blackbox/analyzers/performance_analyzer.py b/envs/monkey_zoo/blackbox/analyzers/performance_analyzer.py
index 92d8a68eb..665ec18af 100644
--- a/envs/monkey_zoo/blackbox/analyzers/performance_analyzer.py
+++ b/envs/monkey_zoo/blackbox/analyzers/performance_analyzer.py
@@ -4,6 +4,7 @@ from typing import Dict
 
 from envs.monkey_zoo.blackbox.analyzers.analyzer import Analyzer
 from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+
 LOGGER = logging.getLogger(__name__)
 
 
diff --git a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py
index 22fa57ac1..e3ecb6eb8 100644
--- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py
+++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py
@@ -1,8 +1,7 @@
-from datetime import timedelta
-from time import sleep
 import json
-
 import logging
+from time import sleep
+
 from bson import json_util
 
 from envs.monkey_zoo.blackbox.island_client.monkey_island_requests import MonkeyIslandRequests
diff --git a/envs/monkey_zoo/blackbox/island_client/monkey_island_requests.py b/envs/monkey_zoo/blackbox/island_client/monkey_island_requests.py
index b56f51612..edcffe674 100644
--- a/envs/monkey_zoo/blackbox/island_client/monkey_island_requests.py
+++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_requests.py
@@ -7,9 +7,9 @@ import functools
 
 from envs.monkey_zoo.blackbox.island_client.supported_reuqest_method import SupportedRequestMethod
 
-# SHA3-512 of '1234567890!@#$%^&*()_nothing_up_my_sleeve_1234567890!@#$%^&*()'
 import logging
 
+# SHA3-512 of '1234567890!@#$%^&*()_nothing_up_my_sleeve_1234567890!@#$%^&*()'
 NO_AUTH_CREDS = '55e97c9dcfd22b8079189ddaeea9bce8125887e3237b800c6176c9afa80d2062' \
                 '8d2c8d0b1538d2208c1444ac66535b764a3d902b35e751df3faec1e477ed3557'
 LOGGER = logging.getLogger(__name__)
diff --git a/envs/monkey_zoo/blackbox/log_handlers/monkey_log.py b/envs/monkey_zoo/blackbox/log_handlers/monkey_log.py
index 091be570a..b7f424a69 100644
--- a/envs/monkey_zoo/blackbox/log_handlers/monkey_log.py
+++ b/envs/monkey_zoo/blackbox/log_handlers/monkey_log.py
@@ -1,6 +1,6 @@
+import logging
 import os
 
-import logging
 from bson import ObjectId
 
 LOGGER = logging.getLogger(__name__)
diff --git a/envs/monkey_zoo/blackbox/log_handlers/test_logs_handler.py b/envs/monkey_zoo/blackbox/log_handlers/test_logs_handler.py
index b54f773e6..bae6a9adc 100644
--- a/envs/monkey_zoo/blackbox/log_handlers/test_logs_handler.py
+++ b/envs/monkey_zoo/blackbox/log_handlers/test_logs_handler.py
@@ -1,8 +1,7 @@
+import logging
 import os
 import shutil
 
-import logging
-
 from envs.monkey_zoo.blackbox.log_handlers.monkey_log_parser import MonkeyLogParser
 from envs.monkey_zoo.blackbox.log_handlers.monkey_logs_downloader import MonkeyLogsDownloader
 
diff --git a/envs/monkey_zoo/blackbox/test_blackbox.py b/envs/monkey_zoo/blackbox/test_blackbox.py
index 631229e36..3c4b5e3df 100644
--- a/envs/monkey_zoo/blackbox/test_blackbox.py
+++ b/envs/monkey_zoo/blackbox/test_blackbox.py
@@ -1,18 +1,16 @@
-import os
 import logging
-
-import pytest
+import os
 from time import sleep
 
-from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
+import pytest
+
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
 from envs.monkey_zoo.blackbox.island_client.island_config_parser import IslandConfigParser
-from envs.monkey_zoo.blackbox.tests.performance.map_generation import MapGenerationTest
-from envs.monkey_zoo.blackbox.tests.performance.report_generation import ReportGenerationTest
+from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
+from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
+from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
 from envs.monkey_zoo.blackbox.tests.performance.telemetry_performance_test import TelemetryPerformanceTest
 from envs.monkey_zoo.blackbox.utils import gcp_machine_handlers
-from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
-from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
 
 DEFAULT_TIMEOUT_SECONDS = 5*60
 MACHINE_BOOTUP_WAIT_SECONDS = 30
diff --git a/envs/monkey_zoo/blackbox/tests/exploitation.py b/envs/monkey_zoo/blackbox/tests/exploitation.py
index e731d8f90..2d55f2294 100644
--- a/envs/monkey_zoo/blackbox/tests/exploitation.py
+++ b/envs/monkey_zoo/blackbox/tests/exploitation.py
@@ -1,9 +1,8 @@
+import logging
 from time import sleep
 
-import logging
-
-from envs.monkey_zoo.blackbox.utils.test_timer import TestTimer
 from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.utils.test_timer import TestTimer
 
 MAX_TIME_FOR_MONKEYS_TO_DIE = 5 * 60
 WAIT_TIME_BETWEEN_REQUESTS = 10
diff --git a/envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py b/envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py
index 5a83f252c..5b8a43227 100644
--- a/envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py
+++ b/envs/monkey_zoo/blackbox/tests/performance/endpoint_performance_test.py
@@ -1,11 +1,10 @@
 import logging
 
-from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
-from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
-from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 from envs.monkey_zoo.blackbox.analyzers.performance_analyzer import PerformanceAnalyzer
+from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
 from envs.monkey_zoo.blackbox.island_client.supported_reuqest_method import SupportedRequestMethod
-
+from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 
 LOGGER = logging.getLogger(__name__)
 
diff --git a/envs/monkey_zoo/blackbox/tests/performance/map_generation.py b/envs/monkey_zoo/blackbox/tests/performance/map_generation.py
index c597907f4..eb95fdc6a 100644
--- a/envs/monkey_zoo/blackbox/tests/performance/map_generation.py
+++ b/envs/monkey_zoo/blackbox/tests/performance/map_generation.py
@@ -1,8 +1,8 @@
 from datetime import timedelta
 
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
-from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 from envs.monkey_zoo.blackbox.tests.performance.performance_test_workflow import PerformanceTestWorkflow
 
 MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
diff --git a/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py b/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py
index 3157140a9..f6cd1dada 100644
--- a/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py
+++ b/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py
@@ -1,7 +1,7 @@
 from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
-from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 from envs.monkey_zoo.blackbox.tests.performance.endpoint_performance_test import EndpointPerformanceTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 
 
 class PerformanceTestWorkflow(BasicTest):
diff --git a/envs/monkey_zoo/blackbox/tests/performance/report_generation.py b/envs/monkey_zoo/blackbox/tests/performance/report_generation.py
index 52fe76288..e204cc29f 100644
--- a/envs/monkey_zoo/blackbox/tests/performance/report_generation.py
+++ b/envs/monkey_zoo/blackbox/tests/performance/report_generation.py
@@ -1,9 +1,9 @@
 from datetime import timedelta
 
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
 from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
 from envs.monkey_zoo.blackbox.tests.performance.performance_test_workflow import PerformanceTestWorkflow
-from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
 
 MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
 MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
diff --git a/envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test.py b/envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test.py
index 0950be4ff..236625288 100644
--- a/envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test.py
+++ b/envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test.py
@@ -1,12 +1,14 @@
+import json
 import logging
 from datetime import timedelta
-import json
 
-from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
-from envs.monkey_zoo.blackbox.tests.performance.utils.telem_parser import TelemParser
+from tqdm import tqdm
+
 from envs.monkey_zoo.blackbox.analyzers.performance_analyzer import PerformanceAnalyzer
-from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
 from envs.monkey_zoo.blackbox.island_client.supported_reuqest_method import SupportedRequestMethod
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.utils.telem_parser import TelemParser
 
 LOGGER = logging.getLogger(__name__)
 
@@ -30,7 +32,7 @@ class TelemetryPerformanceTest:
         LOGGER.info("Telemetries imported successfully.")
         all_telemetries.sort(key=lambda telem: telem['time']['$date'])
         telemetry_parse_times = {}
-        for telemetry in all_telemetries:
+        for telemetry in tqdm(all_telemetries, total=len(all_telemetries), ascii=True, desc="Telemetries sent"):
             telemetry_endpoint = TelemetryPerformanceTest.get_verbose_telemetry_endpoint(telemetry)
             telemetry_parse_times[telemetry_endpoint] = self.get_telemetry_time(telemetry)
         test_config = PerformanceTestConfig(MAX_ALLOWED_SINGLE_TELEM_PARSE_TIME, MAX_ALLOWED_TOTAL_TIME)
diff --git a/envs/monkey_zoo/blackbox/tests/performance/utils/telem_parser.py b/envs/monkey_zoo/blackbox/tests/performance/utils/telem_parser.py
index 06846d17f..049454a32 100644
--- a/envs/monkey_zoo/blackbox/tests/performance/utils/telem_parser.py
+++ b/envs/monkey_zoo/blackbox/tests/performance/utils/telem_parser.py
@@ -1,14 +1,17 @@
-from typing import List, Dict
-from os import listdir, path
 import copy
 import json
+import logging
 import sys
+from os import listdir, path
+from typing import List, Dict
+
+from tqdm import tqdm
 
 from envs.monkey_zoo.blackbox.tests.performance.utils.fake_ip_generator import FakeIpGenerator
 from envs.monkey_zoo.blackbox.tests.performance.utils.fake_monkey import FakeMonkey
 
-
 TELEM_DIR_PATH = './tests/performance/test_telems'
+LOGGER = logging.getLogger(__name__)
 
 
 class TelemParser:
@@ -21,7 +24,7 @@ class TelemParser:
         telems = TelemParser.get_all_telemetries()
         telem_contents = [json.loads(telem['content']) for telem in telems]
         monkeys = self.get_monkeys_from_telems(telem_contents)
-        for i in range(self.multiplier):
+        for i in tqdm(range(self.multiplier), desc="Batch of fabricated telemetries"):
             for monkey in monkeys:
                 monkey.change_fake_data()
             fake_telem_batch = copy.deepcopy(telems)
@@ -31,7 +34,7 @@ class TelemParser:
 
     @staticmethod
     def fabricate_monkeys_in_telems(telems: List[Dict], monkeys: List[FakeMonkey]):
-        for telem in telems:
+        for telem in tqdm(telems, desc="Teletry"):
             for monkey in monkeys:
                 if monkey.on_island:
                     continue
diff --git a/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py b/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py
index 3cb2ad6af..633f406a5 100644
--- a/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py
+++ b/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py
@@ -1,6 +1,5 @@
-import subprocess
-
 import logging
+import subprocess
 
 LOGGER = logging.getLogger(__name__)
 
diff --git a/envs/monkey_zoo/blackbox/utils/json_encoder.py b/envs/monkey_zoo/blackbox/utils/json_encoder.py
index 77be9211a..1642ea222 100644
--- a/envs/monkey_zoo/blackbox/utils/json_encoder.py
+++ b/envs/monkey_zoo/blackbox/utils/json_encoder.py
@@ -1,4 +1,5 @@
 import json
+
 from bson import ObjectId
 
 
diff --git a/monkey/monkey_island/cc/resources/test/utils/telem_store.py b/monkey/monkey_island/cc/resources/test/utils/telem_store.py
index 495a0ff88..49d7ea27b 100644
--- a/monkey/monkey_island/cc/resources/test/utils/telem_store.py
+++ b/monkey/monkey_island/cc/resources/test/utils/telem_store.py
@@ -9,6 +9,7 @@ from monkey_island.cc.models.test_telem import TestTelem
 from monkey_island.cc.services.config import ConfigService
 
 TEST_TELEM_DIR = "./test_telems"
+MAX_SAME_CATEGORY_TELEMS = 10000
 
 
 class TestTelemStore:
@@ -42,11 +43,12 @@ class TestTelemStore:
     @staticmethod
     def get_unique_file_path_for_test_telem(target_dir: str, test_telem: TestTelem):
         telem_filename = TestTelemStore._get_filename_by_test_telem(test_telem)
-        for i in range(100):
+        for i in range(MAX_SAME_CATEGORY_TELEMS):
             potential_filepath = path.join(target_dir, (telem_filename + str(i)))
             if path.exists(potential_filepath):
                 continue
             return potential_filepath
+        raise Exception(f"Too many telemetries of the same category. Max amount {MAX_SAME_CATEGORY_TELEMS}")
 
     @staticmethod
     def _get_filename_by_test_telem(test_telem: TestTelem):
diff --git a/monkey/monkey_island/requirements.txt b/monkey/monkey_island/requirements.txt
index cad53d1c8..39364cbb5 100644
--- a/monkey/monkey_island/requirements.txt
+++ b/monkey/monkey_island/requirements.txt
@@ -24,3 +24,4 @@ requests
 dpath
 ring
 stix2
+tcdm