From faf6da15bb55f7417635dd37cb2a65276e74680c Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Mon, 2 Sep 2019 10:20:52 +0300 Subject: [PATCH] Improved doc, refactored names and added test case for segmentation_utils CR --- monkey/common/network/segmentation_utils.py | 9 +++++++-- .../common/network/segmentation_utils_test.py | 17 ++++++++++++++--- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/monkey/common/network/segmentation_utils.py b/monkey/common/network/segmentation_utils.py index 6569d636b..aeff7f135 100644 --- a/monkey/common/network/segmentation_utils.py +++ b/monkey/common/network/segmentation_utils.py @@ -11,8 +11,13 @@ def get_ip_in_src_and_not_in_dst(ip_addresses, source_subnet, target_subnet): return get_ip_if_in_subnet(ip_addresses, source_subnet) -def get_ip_if_in_subnet(ip_addresses, source_subnet): +def get_ip_if_in_subnet(ip_addresses, subnet): + """ + :param ip_addresses: IP address list. + :param subnet: Subnet to check if one of ip_addresses is in there. This is common.network.network_range.NetworkRange + :return: The first IP in ip_addresses which is in the subnet. + """ for ip_address in ip_addresses: - if source_subnet.is_in_range(ip_address): + if subnet.is_in_range(ip_address): return ip_address return None diff --git a/monkey/common/network/segmentation_utils_test.py b/monkey/common/network/segmentation_utils_test.py index 7ef3e4450..56a560922 100644 --- a/monkey/common/network/segmentation_utils_test.py +++ b/monkey/common/network/segmentation_utils_test.py @@ -8,12 +8,23 @@ class TestSegmentationUtils(IslandTestCase): self.fail_if_not_testing_env() source = CidrRange("1.1.1.0/24") target = CidrRange("2.2.2.0/24") - self.assertIsNone(get_ip_in_src_and_not_in_dst( - [text_type("2.2.2.2")], source, target - )) + + # IP not in both self.assertIsNone(get_ip_in_src_and_not_in_dst( [text_type("3.3.3.3"), text_type("4.4.4.4")], source, target )) + + # IP not in source, in target + self.assertIsNone(get_ip_in_src_and_not_in_dst( + [text_type("2.2.2.2")], source, target + )) + + # IP in source, not in target self.assertIsNotNone(get_ip_in_src_and_not_in_dst( [text_type("8.8.8.8"), text_type("1.1.1.1")], source, target )) + + # IP in both subnets + self.assertIsNone(get_ip_in_src_and_not_in_dst( + [text_type("8.8.8.8"), text_type("1.1.1.1")], source, source + ))