Add fn to check if disabled in config

This commit is contained in:
Shreya 2020-07-14 17:42:49 +05:30
parent cd7cdf846e
commit 50aa068639
1 changed files with 10 additions and 6 deletions

View File

@ -62,16 +62,16 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
Gets the status of a certain attack technique. Gets the status of a certain attack technique.
:return: ScanStatus numeric value :return: ScanStatus numeric value
""" """
if mongo.db.telemetry.find_one({'telem_category': 'attack', if cls._is_disabled_in_config():
'data.status': ScanStatus.USED.value, return ScanStatus.DISABLED.value
'data.technique': cls.tech_id}): elif mongo.db.telemetry.find_one({'telem_category': 'attack',
'data.status': ScanStatus.USED.value,
'data.technique': cls.tech_id}):
return ScanStatus.USED.value return ScanStatus.USED.value
elif mongo.db.telemetry.find_one({'telem_category': 'attack', elif mongo.db.telemetry.find_one({'telem_category': 'attack',
'data.status': ScanStatus.SCANNED.value, 'data.status': ScanStatus.SCANNED.value,
'data.technique': cls.tech_id}): 'data.technique': cls.tech_id}):
return ScanStatus.SCANNED.value return ScanStatus.SCANNED.value
elif not AttackConfig.get_technique_values()[cls.tech_id]:
return ScanStatus.DISABLED.value
else: else:
return ScanStatus.UNSCANNED.value return ScanStatus.UNSCANNED.value
@ -143,6 +143,10 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
@classmethod @classmethod
def _check_status(cls, status): def _check_status(cls, status):
if status == ScanStatus.UNSCANNED.value and not AttackConfig.get_technique_values()[cls.tech_id]: if status == ScanStatus.UNSCANNED.value and cls._is_disabled_in_config():
return ScanStatus.DISABLED.value return ScanStatus.DISABLED.value
return status return status
@classmethod
def _is_disabled_in_config(cls):
return not AttackConfig.get_technique_values()[cls.tech_id]