Rename probability related methods and classes away from TTC since they can be used with any probabilities

Author: andrewbwm

maltoolbox/language/languagegraph.py

@@ -654,7 +654,8 @@ class LanguageGraph():
     """Graph representation of a MAL language"""
     def __init__(self, lang: Optional[dict] = None):
         self.assets: dict = {}
-        self.load_predefined_ttcs()
+        self.predef_ttcs = load_dict_from_yaml_file(
+            'maltoolbox/language/predefined_ttcs.yml')
         if lang is not None:
             self._lang_spec: dict = lang
             self.metadata = {
@@ -696,11 +697,6 @@ def from_mar_archive(cls, mar_archive: str) -> LanguageGraph:
             return LanguageGraph(json.loads(langspec))
 
 
-    def load_predefined_ttcs(self):
-        """Load the predefined ttcs into a dictionary"""
-        self.predef_ttcs = load_dict_from_yaml_file(
-            'maltoolbox/language/predefined_ttcs.yml')
-
     def replace_if_predef_ttc(self, ttc_entry: dict) -> dict:
         """
         If the TTC provided is a predefined name replace it with the

maltoolbox/probs_utils.py

@@ -7,12 +7,12 @@
 
 logger = logging.getLogger(__name__)
 
-class TTCCalculation(Enum):
+class ProbCalculationMethod(Enum):
     SAMPLE = 1
     EXPECTED = 2
 
-def sample_ttc(probs_dict):
-    """Calculate the sampled value from a ttc distribution function
+def sample_prob(probs_dict):
+    """Calculate the sampled value from a probability distribution function
     Arguments:
     probs_dict      - a dictionary containing the probability distribution
                       function that is part of a TTC
@@ -67,8 +67,8 @@ def sample_ttc(probs_dict):
                 f'function encountered {probs_dict["name"]}!')
 
 
-def expected_ttc(probs_dict):
-    """Calculate the expected value from a ttc distribution function
+def expected_prob(probs_dict):
+    """Calculate the expected value from a probability distribution function
     Arguments:
     probs_dict      - a dictionary containing the probability distribution
                       function that is part of a TTC
@@ -122,8 +122,8 @@ def expected_ttc(probs_dict):
                 f'function encountered {probs_dict["name"]}!')
 
 
-def calculate_ttc(probs_dict: dict, method: TTCCalculation) -> float:
-    """Calculate the value from a ttc distribution
+def calculate_prob(probs_dict: dict, method: ProbCalculationMethod) -> float:
+    """Calculate the value from a probability distribution
     Arguments:
     probs_dict      - a dictionary containing the probability distribution
                       corresponding to the TTC
@@ -142,8 +142,8 @@ def calculate_ttc(probs_dict: dict, method: TTCCalculation) -> float:
     match(probs_dict['type']):
         case 'addition' | 'subtraction' | 'multiplication' | \
                 'division' | 'exponentiation':
-            lv = calculate_ttc(probs_dict['lhs'], method)
-            rv = calculate_ttc(probs_dict['rhs'], method)
+            lv = calculate_prob(probs_dict['lhs'], method)
+            rv = calculate_prob(probs_dict['rhs'], method)
             match(probs_dict['type']):
                 case 'addition':
                     return lv + rv
@@ -161,10 +161,10 @@ def calculate_ttc(probs_dict: dict, method: TTCCalculation) -> float:
 
         case 'function':
             match(method):
-                case TTCCalculation.SAMPLE:
-                    return sample_ttc(probs_dict)
-                case TTCCalculation.EXPECTED:
-                    return expected_ttc(probs_dict)
+                case ProbCalculationMethod.SAMPLE:
+                    return sample_prob(probs_dict)
+                case ProbCalculationMethod.EXPECTED:
+                    return expected_prob(probs_dict)
                 case _:
                     raise ValueError('Unknown TTC Calculation method '
                     f'encountered {method}!')

tests/test_probs_utils.py

@@ -4,7 +4,7 @@
 
 from maltoolbox.model import Model
 from maltoolbox.attackgraph.attackgraph import AttackGraph
-from maltoolbox.probs_utils import calculate_ttc, TTCCalculation
+from maltoolbox.probs_utils import calculate_prob, ProbCalculationMethod
 
 def test_probs_utils(model: Model):
     """Test TTC calculation for nodes"""
@@ -24,8 +24,8 @@ def test_probs_utils(model: Model):
 
     for node in attack_graph.nodes.values():
         #TODO: Actually check some of the results
-        calculate_ttc(node.ttc, TTCCalculation.SAMPLE)
+        calculate_prob(node.ttc, ProbCalculationMethod.SAMPLE)
 
     for node in attack_graph.nodes.values():
         #TODO: Actually check some of the results
-        calculate_ttc(node.ttc, TTCCalculation.EXPECTED)
+        calculate_prob(node.ttc, ProbCalculationMethod.EXPECTED)