From 28601d97ed443922230b181f38a4ea28ca43a6ea Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 4 Jan 2021 17:04:31 +0200 Subject: [PATCH] Refactored utility method and added unit test --- monkey/common/network/network_utils.py | 9 --------- monkey/common/utils/code_utils.py | 12 ++---------- monkey/common/utils/test_code_utils.py | 9 +++++++++ monkey/infection_monkey/monkey.py | 2 +- monkey/infection_monkey/network/tools.py | 8 ++++++++ .../system_info/collectors/aws_collector.py | 2 +- .../scoutsuite/data_parsing/rule_parsing.py | 4 ++-- 7 files changed, 23 insertions(+), 23 deletions(-) create mode 100644 monkey/common/utils/test_code_utils.py diff --git a/monkey/common/network/network_utils.py b/monkey/common/network/network_utils.py index 384017f3b..e99d0cf2b 100644 --- a/monkey/common/network/network_utils.py +++ b/monkey/common/network/network_utils.py @@ -1,9 +1,6 @@ import re from urllib.parse import urlparse -from infection_monkey.config import WormConfiguration -from infection_monkey.network.tools import is_running_on_server - def get_host_from_network_location(network_location: str) -> str: """ @@ -21,9 +18,3 @@ def remove_port(url): with_port = f'{parsed.scheme}://{parsed.netloc}' without_port = re.sub(':[0-9]+(?=$|\/)', '', with_port) return without_port - - -def is_running_on_island(): - current_server_without_port = get_host_from_network_location(WormConfiguration.current_server) - running_on_island = is_running_on_server(current_server_without_port) - return running_on_island and WormConfiguration.depth == WormConfiguration.max_depth diff --git a/monkey/common/utils/code_utils.py b/monkey/common/utils/code_utils.py index 32f972a76..d7fad9062 100644 --- a/monkey/common/utils/code_utils.py +++ b/monkey/common/utils/code_utils.py @@ -2,7 +2,7 @@ # noinspection PyPep8Naming import operator from functools import reduce -from typing import List +from typing import List, Union, Any class abstractstatic(staticmethod): @@ -15,13 +15,5 @@ class abstractstatic(staticmethod): __isabstractmethod__ = True -def _get_value_by_path(data, path: List[str]): +def get_dict_value_by_path(data: dict, path: List[str]) -> Any: return reduce(operator.getitem, path, data) - - -def get_object_value_by_path(data_object: object, path: List[str]): - return _get_value_by_path(data_object, path) - - -def get_dict_value_by_path(data_dict: dict, path: List[str]): - return _get_value_by_path(data_dict, path) diff --git a/monkey/common/utils/test_code_utils.py b/monkey/common/utils/test_code_utils.py new file mode 100644 index 000000000..b383b0a33 --- /dev/null +++ b/monkey/common/utils/test_code_utils.py @@ -0,0 +1,9 @@ +import unittest + +from common.utils.code_utils import get_dict_value_by_path + + +class TestCodeUtils(unittest.TestCase): + def test_get_dict_value_by_path(self): + dict_for_test = {'a': {'b': {'c': 'result'}}} + self.assertEqual(get_dict_value_by_path(dict_for_test, ['a', 'b', 'c']), 'result') diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 0e0718f42..e67e2e4d8 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -7,7 +7,7 @@ import time from threading import Thread import infection_monkey.tunnel as tunnel -from common.network.network_utils import is_running_on_island +from infection_monkey.network.tools import is_running_on_island from common.utils.attack_utils import ScanStatus, UsageEnum from common.utils.exceptions import ExploitingVulnerableMachineError, FailedExploitationError from common.version import get_version diff --git a/monkey/infection_monkey/network/tools.py b/monkey/infection_monkey/network/tools.py index 049b30838..b0620abec 100644 --- a/monkey/infection_monkey/network/tools.py +++ b/monkey/infection_monkey/network/tools.py @@ -7,6 +7,8 @@ import subprocess import sys import time +from common.network.network_utils import get_host_from_network_location +from infection_monkey.config import WormConfiguration from infection_monkey.network.info import get_routes, local_ips from infection_monkey.pyinstaller_utils import get_binary_file_path from infection_monkey.utils.environment import is_64bit_python @@ -311,5 +313,11 @@ def get_interface_to_target(dst): return ret[1] +def is_running_on_island(): + current_server_without_port = get_host_from_network_location(WormConfiguration.current_server) + running_on_island = is_running_on_server(current_server_without_port) + return running_on_island and WormConfiguration.depth == WormConfiguration.max_depth + + def is_running_on_server(ip: str) -> bool: return ip in local_ips() diff --git a/monkey/infection_monkey/system_info/collectors/aws_collector.py b/monkey/infection_monkey/system_info/collectors/aws_collector.py index aab70066e..94a7baf2a 100644 --- a/monkey/infection_monkey/system_info/collectors/aws_collector.py +++ b/monkey/infection_monkey/system_info/collectors/aws_collector.py @@ -3,7 +3,7 @@ import logging from common.cloud.aws.aws_instance import AwsInstance from common.cloud.scoutsuite_consts import CloudProviders from common.common_consts.system_info_collectors_names import AWS_COLLECTOR -from common.network.network_utils import is_running_on_island +from infection_monkey.network.tools import is_running_on_island from infection_monkey.system_info.collectors.scoutsuite_collector.scoutsuite_collector import scan_cloud_security from infection_monkey.system_info.system_info_collector import SystemInfoCollector diff --git a/monkey/monkey_island/cc/services/zero_trust/scoutsuite/data_parsing/rule_parsing.py b/monkey/monkey_island/cc/services/zero_trust/scoutsuite/data_parsing/rule_parsing.py index ed2530755..c5855ddd5 100644 --- a/monkey/monkey_island/cc/services/zero_trust/scoutsuite/data_parsing/rule_parsing.py +++ b/monkey/monkey_island/cc/services/zero_trust/scoutsuite/data_parsing/rule_parsing.py @@ -1,4 +1,4 @@ -from common.utils.code_utils import get_object_value_by_path +from common.utils.code_utils import get_dict_value_by_path from common.utils.exceptions import RulePathCreatorNotFound from monkey_island.cc.services.zero_trust.scoutsuite.data_parsing.rule_path_building.rule_path_creators_list import \ RULE_PATH_CREATORS_LIST @@ -9,7 +9,7 @@ class RuleParser: @staticmethod def get_rule_data(scoutsuite_data, rule_name): rule_path = RuleParser.get_rule_path(rule_name) - return get_object_value_by_path(scoutsuite_data, rule_path) + return get_dict_value_by_path(data=scoutsuite_data, path=rule_path) @staticmethod def get_rule_path(rule_name):