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 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:<br/>"
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])}<br/>"
f"{', '.join(config_schema_per_attack_technique[cls.tech_id][config_type])}<br/>"
)
return config_options

View File

@ -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,
)