Add optional parameter 'configuration' in _check_status()
This commit is contained in:
parent
6512d5de11
commit
9217581e1c
|
@ -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]
|
||||||
|
|
|
@ -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)
|
||||||
|
|
Loading…
Reference in New Issue