Add optional parameter 'configuration' in _check_status()

This commit is contained in:
Shreya 2020-07-18 23:43:47 +05:30
parent 6512d5de11
commit 9217581e1c
2 changed files with 36 additions and 29 deletions

View File

@ -63,7 +63,7 @@ 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 cls._is_disabled_in_config(): if not cls._is_enabled_in_config():
return ScanStatus.DISABLED.value return ScanStatus.DISABLED.value
elif mongo.db.telemetry.find_one({'telem_category': 'attack', elif mongo.db.telemetry.find_one({'telem_category': 'attack',
'data.status': ScanStatus.USED.value, 'data.status': ScanStatus.USED.value,
@ -143,11 +143,12 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
return {} return {}
@classmethod @classmethod
def _check_status(cls, status): def _check_status(cls, status, *args):
if status == ScanStatus.UNSCANNED.value and cls._is_disabled_in_config(): enabled_in_config = args[0] if args else cls._is_enabled_in_config()
if status == ScanStatus.UNSCANNED.value and not enabled_in_config:
return ScanStatus.DISABLED.value return ScanStatus.DISABLED.value
return status return status
@classmethod @classmethod
def _is_disabled_in_config(cls): def _is_enabled_in_config(cls):
return not AttackConfig.get_technique_values()[cls.tech_id] return AttackConfig.get_technique_values()[cls.tech_id]

View File

@ -13,27 +13,27 @@ class T9999(PostBreachTechnique):
pba_names = ["PBA Name"] pba_names = ["PBA Name"]
config =\ # config =\
{ # {
'category': { # 'category': {
'link': '', # 'link': '',
'properties': { # 'properties': {
'T9999': { # 'T9999': {
'description': '', # 'description': '',
'link': '', # 'link': '',
'necessary': False, # 'necessary': False,
'title': '', # 'title': '',
'type': 'bool', # 'type': 'bool',
'value': False # this field denotes whether technique is enabled or disabled in config # 'value': False # this field denotes whether technique is enabled or disabled in config
} # }
} # }
} # }
} # }
def set_config(value: bool): # def set_config(value: bool):
config['category']['properties']['T9999']['value'] = value # config['category']['properties']['T9999']['value'] = value
return config # return config
param_list = [ param_list = [
@ -45,7 +45,8 @@ param_list = [
'status': 0, 'status': 0,
'title': 'Unscanned technique' 'title': 'Unscanned technique'
}, },
set_config(True), # configuration # set_config(True), # configuration
True, # configuration
ScanStatus.UNSCANNED.value # expected status ScanStatus.UNSCANNED.value # expected status
), ),
@ -57,7 +58,8 @@ param_list = [
'status': 1, 'status': 1,
'title': 'Scanned technique' 'title': 'Scanned technique'
}, },
set_config(True), # configuration # set_config(True), # configuration
True, # configuration
ScanStatus.SCANNED.value # expected status ScanStatus.SCANNED.value # expected status
), ),
@ -69,7 +71,8 @@ param_list = [
'status': 2, 'status': 2,
'title': 'Used technique' 'title': 'Used technique'
}, },
set_config(True), # configuration # set_config(True), # configuration
True, # configuration
ScanStatus.USED.value # expected status ScanStatus.USED.value # expected status
), ),
@ -81,7 +84,8 @@ param_list = [
'status': 0, 'status': 0,
'title': 'Disabled technique' 'title': 'Disabled technique'
}, },
set_config(False), # configuration # set_config(False), # configuration
False, # configuration
ScanStatus.DISABLED.value # expected status ScanStatus.DISABLED.value # expected status
) )
] ]
@ -91,4 +95,6 @@ class TestAttackTechnique(TestCase):
def test__check_status(self): def test__check_status(self):
for telem_data, config, expected_status in param_list: for telem_data, config, expected_status in param_list:
with self.subTest(msg=f"Checking if correct status is returned (status: {telem_data['message']})"): with self.subTest(msg=f"Checking if correct status is returned (status: {telem_data['message']})"):
self.assertEqual(T9999._check_status(telem_data['status']), expected_status) # self.assertEqual(T9999._check_status(telem_data['status']), expected_status)
self.assertEqual(T9999._check_status(telem_data['status'], config),
expected_status)