From 4959a09c089685ff476fdad2bf108487550c7786 Mon Sep 17 00:00:00 2001 From: Brent Haub Date: Tue, 14 Jan 2025 15:03:10 +0000 Subject: [PATCH 1/2] Update typeChecker in exceptions.py to support a list of allowed types --- mitreattack/navlayers/core/exceptions.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mitreattack/navlayers/core/exceptions.py b/mitreattack/navlayers/core/exceptions.py index c7125b16..d8deaad5 100644 --- a/mitreattack/navlayers/core/exceptions.py +++ b/mitreattack/navlayers/core/exceptions.py @@ -55,13 +55,18 @@ def typeChecker(caller, testee, desired_type, field): :param caller: the entity that called this function (used for error messages) :param testee: the element to test - :param desired_type: the type the element should be + :param desired_type: the type the element should be or a list of + allowed types :param field: what the element is to be used as (used for error messages) :raises BadType: error denoting the testee element is not of the correct type """ - if not isinstance(testee, desired_type): + if isinstance(desired_type, list): + if not any(isinstance(testee, t) for t in desired_type): + handler(caller, f"{testee} [{field}] is not one of {str(desired_type)}") + raise BadType + elif not isinstance(testee, desired_type): handler(caller, f"{testee} [{field}] is not a {str(desired_type)}") raise BadType From 90ed385b4eeed1c383fc94fb00c5949d3a473666 Mon Sep 17 00:00:00 2001 From: Brent Haub Date: Tue, 14 Jan 2025 15:03:23 +0000 Subject: [PATCH 2/2] Update technique.py to utilize list of allowed types for score attribute type check --- mitreattack/navlayers/core/technique.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/mitreattack/navlayers/core/technique.py b/mitreattack/navlayers/core/technique.py index 21d2a4a7..c66d5050 100644 --- a/mitreattack/navlayers/core/technique.py +++ b/mitreattack/navlayers/core/technique.py @@ -94,12 +94,8 @@ def score(self): @score.setter def score(self, score): """Setter for score.""" - try: - typeChecker(type(self).__name__, score, int, "score") - self.__score = score - except BadType: - typeChecker(type(self).__name__, score, float, "score") - self.__score = int(score) + typeChecker(type(self).__name__, score, [int, float], "score") + self.__score = score @property def color(self):