island, tests: Make config_schema_per_attack_technique a class variable instead of generating it every time

This commit is contained in:
Shreya Malviya 2021-09-27 20:29:30 +05:30
parent afedde8c05
commit 0804cecb64
2 changed files with 23 additions and 15 deletions

View File

@ -1,14 +1,15 @@
import abc import abc
import logging import logging
from typing import List from typing import Dict, List
from common.utils.attack_utils import ScanStatus from common.utils.attack_utils import ScanStatus
from common.utils.code_utils import abstractstatic from common.utils.code_utils import abstractstatic
from monkey_island.cc.database import mongo from monkey_island.cc.database import mongo
from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations 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.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 ( 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__) logger = logging.getLogger(__name__)
@ -108,6 +109,8 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
""" """
return {"message": cls.get_message_by_status(status), "status": status} return {"message": cls.get_message_by_status(status), "status": status}
config_schema_per_attack_technique = None
@classmethod @classmethod
def get_message_by_status(cls, status): def get_message_by_status(cls, status):
""" """
@ -118,7 +121,13 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
if status == ScanStatus.DISABLED.value: if status == ScanStatus.DISABLED.value:
return disabled_msg return disabled_msg
if status == ScanStatus.UNSCANNED.value: 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 return unscanned_msg
elif status == ScanStatus.SCANNED.value: elif status == ScanStatus.SCANNED.value:
return cls.scanned_msg return cls.scanned_msg
@ -126,14 +135,16 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
return cls.used_msg return cls.used_msg
@classmethod @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 = [] reasons = []
if len(cls.relevant_systems) == 1: if len(cls.relevant_systems) == 1:
reasons.append(f"- The Monkey did not run on any {cls.relevant_systems[0]} systems.") 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( reasons.append(
"- The following configuration options were disabled:<br/>" "- The following configuration options were disabled:<br/>"
f"{cls._get_relevant_config_values()}" f"{cls._get_relevant_config_values(config_schema_per_attack_technique)}"
) )
if reasons: if reasons:
@ -146,12 +157,12 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
return unscanned_msg return unscanned_msg
@classmethod @classmethod
def _get_relevant_config_values(cls): def _get_relevant_config_values(cls, config_schema_per_attack_technique: Dict):
config_options = "" 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 += ( config_options += (
f"- {config_type}" f"- {config_type}"
f"{', '.join(CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE[cls.tech_id][config_type])}<br/>" f"{', '.join(config_schema_per_attack_technique[cls.tech_id][config_type])}<br/>"
) )
return config_options return config_options

View File

@ -21,13 +21,10 @@ FAKE_CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE = {
@pytest.fixture(scope="function", autouse=True) @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( monkeypatch.setattr(
( ("monkey_island.cc.services.attack.technique_reports." "__init__.SCHEMA"),
"monkey_island.cc.services.attack.technique_reports." fake_schema,
"__init__.CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE"
),
FAKE_CONFIG_SCHEMA_PER_ATTACK_TECHNIQUE,
) )