From 0804cecb64e8b2fa316eacb54d1dfcafe156fbcb Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 27 Sep 2021 20:29:30 +0530 Subject: [PATCH] island, tests: Make config_schema_per_attack_technique a class variable instead of generating it every time --- .../attack/technique_reports/__init__.py | 29 +++++++++++++------ .../test_technique_reports.py | 9 ++---- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py b/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py index 88187c410..6d3186e73 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py @@ -1,14 +1,15 @@ import abc import logging -from typing import List +from typing import Dict, List from common.utils.attack_utils import ScanStatus from common.utils.code_utils import abstractstatic from monkey_island.cc.database import mongo from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations from monkey_island.cc.services.attack.attack_config import AttackConfig +from monkey_island.cc.services.config_schema.config_schema import SCHEMA from monkey_island.cc.services.config_schema.config_schema_per_attack_technique import ( - CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE, + get_config_schema_per_attack_technique, ) logger = logging.getLogger(__name__) @@ -108,6 +109,8 @@ class AttackTechnique(object, metaclass=abc.ABCMeta): """ return {"message": cls.get_message_by_status(status), "status": status} + config_schema_per_attack_technique = None + @classmethod def get_message_by_status(cls, status): """ @@ -118,7 +121,13 @@ class AttackTechnique(object, metaclass=abc.ABCMeta): if status == ScanStatus.DISABLED.value: return disabled_msg if status == ScanStatus.UNSCANNED.value: - unscanned_msg = cls._get_unscanned_msg_with_reasons(cls.unscanned_msg) + if not cls.config_schema_per_attack_technique: + cls.config_schema_per_attack_technique = get_config_schema_per_attack_technique( + SCHEMA + ) + unscanned_msg = cls._get_unscanned_msg_with_reasons( + cls.unscanned_msg, cls.config_schema_per_attack_technique + ) return unscanned_msg elif status == ScanStatus.SCANNED.value: return cls.scanned_msg @@ -126,14 +135,16 @@ class AttackTechnique(object, metaclass=abc.ABCMeta): return cls.used_msg @classmethod - def _get_unscanned_msg_with_reasons(cls, unscanned_msg): + def _get_unscanned_msg_with_reasons( + cls, unscanned_msg: str, config_schema_per_attack_technique: Dict + ): reasons = [] if len(cls.relevant_systems) == 1: reasons.append(f"- The Monkey did not run on any {cls.relevant_systems[0]} systems.") - if cls.tech_id in CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE: + if cls.tech_id in config_schema_per_attack_technique: reasons.append( "- The following configuration options were disabled:
" - f"{cls._get_relevant_config_values()}" + f"{cls._get_relevant_config_values(config_schema_per_attack_technique)}" ) if reasons: @@ -146,12 +157,12 @@ class AttackTechnique(object, metaclass=abc.ABCMeta): return unscanned_msg @classmethod - def _get_relevant_config_values(cls): + def _get_relevant_config_values(cls, config_schema_per_attack_technique: Dict): config_options = "" - for config_type in CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE[cls.tech_id]: + for config_type in config_schema_per_attack_technique[cls.tech_id]: config_options += ( f"- {config_type} — " - f"{', '.join(CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE[cls.tech_id][config_type])}
" + f"{', '.join(config_schema_per_attack_technique[cls.tech_id][config_type])}
" ) return config_options diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/attack/technique_reports/test_technique_reports.py b/monkey/tests/unit_tests/monkey_island/cc/services/attack/technique_reports/test_technique_reports.py index acd855e71..40a627e3c 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/attack/technique_reports/test_technique_reports.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/attack/technique_reports/test_technique_reports.py @@ -21,13 +21,10 @@ FAKE_CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE = { @pytest.fixture(scope="function", autouse=True) -def mock_config_schema_per_attack_technique(monkeypatch): +def mock_config_schema_per_attack_technique(monkeypatch, fake_schema): monkeypatch.setattr( - ( - "monkey_island.cc.services.attack.technique_reports." - "__init__.CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE" - ), - FAKE_CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE, + ("monkey_island.cc.services.attack.technique_reports." "__init__.SCHEMA"), + fake_schema, )