From 9aff40d974524e0502da1528a8906591729f9514 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Thu, 7 May 2020 11:20:08 +0300
Subject: [PATCH] Fake IP's don't get out of range, fixed typo in requirements
 and added logging to telem parsing

---
 .../performance/utils/fake_ip_generator.py    | 20 ++++++++++++++++---
 .../utils/test_fake_ip_generator.py           | 18 +++++++++++++++++
 .../cc/resources/test/utils/telem_store.py    |  8 ++++++++
 monkey/monkey_island/requirements.txt         |  2 +-
 4 files changed, 44 insertions(+), 4 deletions(-)
 create mode 100644 envs/monkey_zoo/blackbox/tests/performance/utils/test_fake_ip_generator.py

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