Add ttc sampling and expected value to AttackGraphNode

Author: mrkickling

maltoolbox/attackgraph/node.py

@@ -6,11 +6,13 @@
 import copy
 from functools import cached_property
 from typing import TYPE_CHECKING
+import numpy as np
+import math
 
 if TYPE_CHECKING:
     from typing import Any, Optional
     from . import Attacker
-    from ..language import LanguageGraphAttackStep, Detector
+    from ..language import LanguageGraphAttackStep
     from ..model import ModelAsset
 
 class AttackGraphNode:
@@ -181,6 +183,101 @@ def is_available_defense(self) -> bool:
             'suppress' not in self.tags and \
             self.defense_status != 1.0
 
+    def ttc_sample(self) -> float:
+        """Sample a value from ttc distribution for a node
+
+        TTC Distributions in MAL:
+        https://mal-lang.org/mal-langspec/apidocs/org.mal_lang.langspec/org/mal_lang/langspec/ttc/TtcDistribution.html
+        """
+
+        def sample(exponential: float, bernoulli=1.0):
+            """
+            Generate a random sample for the given distributions.
+            If bernoulli distribution is not given, the sample will
+            simply be from an exponential.
+
+            If the Bernoulli trial fails (0), return infinity (impossible).
+            If the Bernoulli trial succeeds (1), return sample from
+            exponential distribution.
+            """
+
+            # If bernoulli is set to 1, the sample is just exponential
+            if np.random.choice([0, 1], p=[1 - bernoulli, bernoulli]):
+                return np.random.exponential(scale=1 / exponential)
+            return math.inf
+
+        if self.type == "defense":
+            # Defenses have no ttc
+            return 0
+
+        distribution = self.ttc.get('name')
+        s = math.nan
+        if distribution == "EasyAndCertain":
+            s = sample(exponential=1)
+        elif distribution == "EasyAndUncertain":
+            s = sample(exponential=1, bernoulli=0.5)
+        elif distribution == "HardAndCertain":
+            s = sample(exponential=0.1)
+        elif distribution == "HardAndUncertain":
+            s = sample(exponential=0.1, bernoulli=0.5)
+        elif distribution == "VeryHardAndCertain":
+            s = sample(exponential=0.01)
+        elif distribution == "VeryHardAndUncertain":
+            s = sample(exponential=0.01, bernoulli=0.5)
+        elif distribution == "Exponential":
+            scale = float(self.ttc['arguments'][0])
+            s = sample(exponential=scale)
+        else:
+            raise ValueError(f"Unknown TTC distribution: {distribution}")
+
+        return s
+
+    def ttc_expected_value(self) -> float:
+        """Returns the expected value of the ttc distribution for a node.
+
+        TTC Distributions in MAL:
+        https://mal-lang.org/mal-langspec/apidocs/org.mal_lang.langspec/org/mal_lang/langspec/ttc/TtcDistribution.html
+        """
+
+        def expected_value(exponential: float, bernoulli=1.0):
+            """Compute expected value for given distributions.
+
+            If bernoulli distribution is 0, the expected value is infinite.
+            Otherwise, the expected value is the expected value of the
+            exponential distribution divided by the probability of the
+            bernoulli distribution.
+            """
+            if bernoulli == 0:
+                # If Bernoulli always blocks, expectation is infinite
+                return math.inf
+
+            # Conditional expectation
+            return (1 / exponential) / bernoulli
+
+        if self.type == "defense":
+            # Defenses have no ttc
+            return 0
+
+        distribution = self.ttc["name"]
+        e = math.nan
+        if distribution == "EasyAndCertain":
+            e = expected_value(exponential=1.0)
+        elif distribution == "EasyAndUncertain":
+            e = expected_value(exponential=1.0, bernoulli=0.5)
+        elif distribution == "HardAndCertain":
+            e = expected_value(exponential=0.1)
+        elif distribution == "HardAndUncertain":
+            e = expected_value(exponential=0.1, bernoulli=0.5)
+        elif distribution == "VeryHardAndCertain":
+            e = expected_value(exponential=0.01)
+        elif distribution == "VeryHardAndUncertain":
+            e = expected_value(exponential=0.01, bernoulli=0.5)
+        elif distribution == "Exponential":
+            scale = float(self.ttc["arguments"][0])
+            e = expected_value(exponential=scale)
+        else:
+            raise ValueError(f"Unknown TTC distribution: {distribution}")
+        return e
 
     @property
     def full_name(self) -> str:

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
   "antlr4-python3-runtime",
   "docopt",
   "PyYAML",
+  "numpy",
 ]
 license = {text = "Apache Software License"}
 keywords = ["mal"]

tests/attackgraph/test_node.py

@@ -1,5 +1,6 @@
 """Unit tests for AttackGraphNode functionality"""
 
+from maltoolbox.model import Model
 from maltoolbox.attackgraph.node import AttackGraphNode
 from maltoolbox.attackgraph.attacker import Attacker
 from maltoolbox.attackgraph.attackgraph import AttackGraph
@@ -83,3 +84,72 @@ def test_attackgraphnode(dummy_lang_graph: LanguageGraph):
     # Node 1 is not a defense
     assert not node1.is_available_defense()
     assert not node1.is_enabled_defense()
+
+
+def test_ttc_node(model: Model):
+    """Test TTC calculation for nodes"""
+
+    app = model.add_asset('Application')
+    creds = model.add_asset('Credentials')
+    user = model.add_asset('User')
+    identity = model.add_asset('Identity')
+    vuln = model.add_asset('SoftwareVulnerability')
+
+    identity.add_associated_assets('credentials', {creds})
+    user.add_associated_assets('userIds', {identity})
+    app.add_associated_assets('highPrivAppIAMs', {identity})
+    app.add_associated_assets('vulnerabilities', {vuln})
+
+    attack_graph = AttackGraph(model.lang_graph, model)
+
+    nodes_with_ttc = [
+        node for node in attack_graph.nodes.values()
+        if node.ttc is not None and node.ttc['name'] != 'Disabled'
+    ]
+
+    assert nodes_with_ttc, "No nodes with ttc set"
+    for node in nodes_with_ttc:
+        # Make sure we can sample and that returned
+        # values are non-negative numbers
+        sample = node.ttc_sample()
+        expected = node.ttc_expected_value()
+        assert sample >= 0, "Sampled TTC is negative"
+        assert expected >= 0, "Expected TTC is negative"
+
+    easy_and_certain_nodes = [
+        node for node in nodes_with_ttc
+        if node.ttc['name'] == 'EasyAndCertain'
+    ]
+    assert easy_and_certain_nodes
+    for node in easy_and_certain_nodes:
+        assert node.ttc_expected_value() == 1.0
+
+
+    # TODO EasyAndUncertain does not exist in coreLang
+    # TODO HardAndCertain does not exist in coreLang
+
+    hard_and_uncertain_nodes = [
+        node for node in nodes_with_ttc
+        if node.ttc['name'] == 'HardAndUncertain'
+    ]
+    assert hard_and_uncertain_nodes
+    for node in hard_and_uncertain_nodes:
+        assert node.ttc_expected_value() == 20.0
+
+    # TODO: VeryHardAndCertain does not exist in coreLang
+
+    very_hard_and_uncertain_nodes = [
+        node for node in nodes_with_ttc
+        if node.ttc['name'] == 'VeryHardAndUncertain'
+    ]
+    assert very_hard_and_uncertain_nodes
+    for node in very_hard_and_uncertain_nodes:
+        assert node.ttc_expected_value() == 200.0
+
+    exponential_nodes = [
+        node for node in nodes_with_ttc
+        if node.ttc['name'] == 'Exponential'
+    ]
+    assert exponential_nodes
+    for node in exponential_nodes:
+        assert node.ttc_expected_value() == (1/node.ttc["arguments"][0])