diff --git a/envs/monkey_zoo/blackbox/tests/performance/utils/fake_ip_generator.py b/envs/monkey_zoo/blackbox/tests/performance/utils/fake_ip_generator.py index 8abe53453..90422f9a0 100644 --- a/envs/monkey_zoo/blackbox/tests/performance/utils/fake_ip_generator.py +++ b/envs/monkey_zoo/blackbox/tests/performance/utils/fake_ip_generator.py @@ -1,11 +1,25 @@ +from typing import List + + class FakeIpGenerator: def __init__(self): self.fake_ip_parts = [1, 1, 1, 1] - def generate_fake_ips_for_real_ips(self, real_ips): - self.fake_ip_parts[2] += 1 + def generate_fake_ips_for_real_ips(self, real_ips: List[str]) -> List[str]: fake_ips = [] for i in range(len(real_ips)): fake_ips.append('.'.join(str(part) for part in self.fake_ip_parts)) - self.fake_ip_parts[3] += 1 + self.increment_ip() return fake_ips + + def increment_ip(self): + self.fake_ip_parts[3] += 1 + self.try_fix_ip_range() + + def try_fix_ip_range(self): + for i in range(len(self.fake_ip_parts)): + if self.fake_ip_parts[i] > 256: + if i-1 < 0: + raise Exception("Fake IP's out of range.") + self.fake_ip_parts[i-1] += 1 + self.fake_ip_parts[i] = 1 diff --git a/envs/monkey_zoo/blackbox/tests/performance/utils/test_fake_ip_generator.py b/envs/monkey_zoo/blackbox/tests/performance/utils/test_fake_ip_generator.py new file mode 100644 index 000000000..96609a8a9 --- /dev/null +++ b/envs/monkey_zoo/blackbox/tests/performance/utils/test_fake_ip_generator.py @@ -0,0 +1,18 @@ +from unittest import TestCase + +from envs.monkey_zoo.blackbox.tests.performance.utils.fake_ip_generator import FakeIpGenerator + + +class TestFakeIpGenerator(TestCase): + + def test_fake_ip_generation(self): + fake_ip_gen = FakeIpGenerator() + self.assertListEqual([1, 1, 1, 1], fake_ip_gen.fake_ip_parts) + for i in range(256): + fake_ip_gen.generate_fake_ips_for_real_ips(['1.1.1.1']) + self.assertListEqual(['1.1.2.1'], fake_ip_gen.generate_fake_ips_for_real_ips(['1.1.1.1'])) + fake_ip_gen.fake_ip_parts = [256, 256, 255, 256] + self.assertListEqual(['256.256.255.256', '256.256.256.1'], + fake_ip_gen.generate_fake_ips_for_real_ips(['1.1.1.1', '1.1.1.2'])) + fake_ip_gen.fake_ip_parts = [256, 256, 256, 256] + self.assertRaises(Exception, fake_ip_gen.generate_fake_ips_for_real_ips(['1.1.1.1'])) 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 49d7ea27b..031a0d02a 100644 --- a/monkey/monkey_island/cc/resources/test/utils/telem_store.py +++ b/monkey/monkey_island/cc/resources/test/utils/telem_store.py @@ -1,3 +1,4 @@ +import logging from functools import wraps from os import mkdir, path import shutil @@ -12,6 +13,10 @@ TEST_TELEM_DIR = "./test_telems" MAX_SAME_CATEGORY_TELEMS = 10000 +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + class TestTelemStore: @staticmethod @@ -31,14 +36,17 @@ class TestTelemStore: @staticmethod def export_test_telems(): + logger.info(f"Exporting all telemetries to {TEST_TELEM_DIR}") try: mkdir(TEST_TELEM_DIR) except FileExistsError: + logger.info("Deleting all previous telemetries.") shutil.rmtree(TEST_TELEM_DIR) mkdir(TEST_TELEM_DIR) for test_telem in TestTelem.objects(): with open(TestTelemStore.get_unique_file_path_for_test_telem(TEST_TELEM_DIR, test_telem), 'w') as file: file.write(test_telem.to_json()) + logger.info("Telemetries exported!") @staticmethod def get_unique_file_path_for_test_telem(target_dir: str, test_telem: TestTelem): diff --git a/monkey/monkey_island/requirements.txt b/monkey/monkey_island/requirements.txt index 39364cbb5..ce376ab4c 100644 --- a/monkey/monkey_island/requirements.txt +++ b/monkey/monkey_island/requirements.txt @@ -24,4 +24,4 @@ requests dpath ring stix2 -tcdm +tqdm