mal-toolbox 0.1.11__py3-none-any.whl → 0.2.0__py3-none-any.whl

maltoolbox/attackgraph/attackgraph.py

@@ -6,193 +6,32 @@ import copy
 import logging
 import json
 
+from itertools import chain
 from typing import TYPE_CHECKING
 
 from .node import AttackGraphNode
 from .attacker import Attacker
-from ..exceptions import AttackGraphStepExpressionError
+from ..exceptions import AttackGraphStepExpressionError, AttackGraphException
+from ..exceptions import LanguageGraphException
 from ..model import Model
-from ..exceptions import AttackGraphException
+from ..language import (LanguageGraph, ExpressionsChain,
+    disaggregate_attack_step_full_name)
 from ..file_utils import (
     load_dict_from_json_file,
     load_dict_from_yaml_file,
     save_dict_to_file
 )
 
+
 if TYPE_CHECKING:
     from typing import Any, Optional
-    from ..language import LanguageGraph
 
 logger = logging.getLogger(__name__)
 
-# TODO see if (part of) this can be incorporated into the LanguageGraph, so that
-# the LanguageGraph's _lang_spec private property does not need to be accessed
-def _process_step_expression(
-        lang_graph: LanguageGraph,
-        model: Model,
-        target_assets: list[Any],
-        step_expression: dict[str, Any]
-    ) -> tuple[list, Optional[str]]:
-    """
-    Recursively process an attack step expression.
-
-    Arguments:
-    lang_graph      - a language graph representing the MAL language
-                      specification
-    model           - a maltoolbox.model.Model instance from which the attack
-                      graph was generated
-    target_assets   - the list of assets that this step expression should apply
-                      to. Initially it will contain the asset to which the
-                      attack step belongs
-    step_expression - a dictionary containing the step expression
-
-    Return:
-    A tuple pair containing a list of all of the target assets and the name of
-    the attack step.
-    """
-
-    if logger.isEnabledFor(logging.DEBUG):
-        # Avoid running json.dumps when not in debug
-        logger.debug(
-            'Processing Step Expression:\n%s',
-            json.dumps(step_expression, indent = 2)
-        )
-
-    match (step_expression['type']):
-        case 'attackStep':
-            # The attack step expression just adds the name of the attack
-            # step. All other step expressions only modify the target assets.
-            return (target_assets, step_expression['name'])
-
-        case 'union' | 'intersection' | 'difference':
-            # The set operators are used to combine the left hand and right
-            # hand targets accordingly.
-            lh_targets, lh_attack_steps = _process_step_expression(
-                lang_graph, model, target_assets, step_expression['lhs'])
-            rh_targets, rh_attack_steps = _process_step_expression(
-                lang_graph, model, target_assets, step_expression['rhs'])
-
-            new_target_assets = []
-            match (step_expression['type']):
-                case 'union':
-                    new_target_assets = lh_targets
-                    for ag_node in rh_targets:
-                        if next((lnode for lnode in new_target_assets \
-                            if lnode.id != ag_node.id), None):
-                            new_target_assets.append(ag_node)
-
-                case 'intersection':
-                    for ag_node in rh_targets:
-                        if next((lnode for lnode in lh_targets \
-                            if lnode.id == ag_node.id), None):
-                            new_target_assets.append(ag_node)
-
-                case 'difference':
-                    new_target_assets = lh_targets
-                    for ag_node in lh_targets:
-                        if next((rnode for rnode in rh_targets \
-                            if rnode.id != ag_node.id), None):
-                            new_target_assets.remove(ag_node)
-
-            return (new_target_assets, None)
-
-        case 'variable':
-            # Fetch the step expression associated with the variable from
-            # the language specification and resolve that.
-            for target_asset in target_assets:
-                if (hasattr(target_asset, 'type')):
-                    # TODO how can this info be accessed in the lang_graph
-                    # directly without going through the private method?
-                    variable_step_expr = lang_graph._get_variable_for_asset_type_by_name(
-                        target_asset.type, step_expression['name'])
-                    return _process_step_expression(
-                        lang_graph, model, target_assets, variable_step_expr)
-
-                else:
-                    logger.error(
-                        'Requested variable from non-asset target node:'
-                        '%s which cannot be resolved.', target_asset
-                    )
-            return ([], None)
-
-        case 'field':
-            # Change the target assets from the current ones to the associated
-            # assets given the specified field name.
-            new_target_assets = []
-            for target_asset in target_assets:
-                new_target_assets.extend(model.\
-                    get_associated_assets_by_field_name(target_asset,
-                        step_expression['name']))
-            return (new_target_assets, None)
-
-        case 'transitive':
-            # The transitive expression is very similar to the field
-            # expression, but it proceeds recursively until no target is
-            # found and it and it sets the new targets to the entire list
-            # of assets identified during the entire transitive recursion.
-            new_target_assets = []
-            for target_asset in target_assets:
-                new_target_assets.extend(model.\
-                    get_associated_assets_by_field_name(target_asset,
-                        step_expression['stepExpression']['name']))
-            if new_target_assets:
-                (additional_assets, _) = _process_step_expression(
-                    lang_graph, model, new_target_assets, step_expression)
-                new_target_assets.extend(additional_assets)
-                return (new_target_assets, None)
-            else:
-                return ([], None)
-
-        case 'subType':
-            new_target_assets = []
-            for target_asset in target_assets:
-                (assets, _) = _process_step_expression(
-                    lang_graph, model, target_assets,
-                    step_expression['stepExpression'])
-                new_target_assets.extend(assets)
-
-            selected_new_target_assets = []
-            for asset in new_target_assets:
-                lang_graph_asset = lang_graph.get_asset_by_name(
-                    asset.type
-                )
-                if not lang_graph_asset:
-                    raise LookupError(
-                        f'Failed to find asset \"{asset.type}\" in the '
-                        'language graph.'
-                    )
-                lang_graph_subtype_asset = lang_graph.get_asset_by_name(
-                    step_expression['subType']
-                )
-                if not lang_graph_subtype_asset:
-                    raise LookupError(
-                        'Failed to find asset '
-                        f'\"{step_expression["subType"]}\" in the '
-                        'language graph.'
-                    )
-                if lang_graph_asset.is_subasset_of(lang_graph_subtype_asset):
-                    selected_new_target_assets.append(asset)
-
-            return (selected_new_target_assets, None)
-
-        case 'collect':
-            # Apply the right hand step expression to left hand step
-            # expression target assets.
-            lh_targets, _ = _process_step_expression(
-                lang_graph, model, target_assets, step_expression['lhs'])
-            return _process_step_expression(lang_graph, model, lh_targets,
-                step_expression['rhs'])
-
-
-        case _:
-            logger.error(
-                'Unknown attack step type: %s', step_expression["type"]
-            )
-            return ([], None)
 
 class AttackGraph():
     """Graph representation of attack steps"""
-    def __init__(self, lang_graph = None, model: Optional[Model] = None):
+    def __init__(self, lang_graph, model: Optional[Model] = None):
         self.nodes: list[AttackGraphNode] = []
         self.attackers: list[Attacker] = []
         # Dictionaries used in optimization to get nodes and attackers by id
@@ -205,7 +44,7 @@ class AttackGraph():
         self.lang_graph = lang_graph
         self.next_node_id = 0
         self.next_attacker_id = 0
-        if self.model is not None and self.lang_graph is not None:
+        if self.model is not None:
             self._generate_graph()
 
     def __repr__(self) -> str:
@@ -220,6 +59,9 @@ class AttackGraph():
                 ag_node.to_dict()
         for attacker in self.attackers:
             serialized_attackers[attacker.name] = attacker.to_dict()
+        logger.debug('Serialized %d attack steps and %d attackers.' %
+            (len(self.nodes), len(self.attackers))
+        )
         return {
             'attack_steps': serialized_attack_steps,
             'attackers': serialized_attackers,
@@ -280,6 +122,7 @@ class AttackGraph():
     def _from_dict(
             cls,
             serialized_object: dict,
+            lang_graph: LanguageGraph,
             model: Optional[Model]=None
         ) -> AttackGraph:
         """Create AttackGraph from dict
@@ -288,13 +131,13 @@ class AttackGraph():
         model               - Optional Model to add connections to
         """
 
-        attack_graph = AttackGraph()
+        attack_graph = AttackGraph(lang_graph)
         attack_graph.model = model
         serialized_attack_steps = serialized_object['attack_steps']
         serialized_attackers = serialized_object['attackers']
 
         # Create all of the nodes in the imported attack graph.
-        for node_full_name, node_dict in serialized_attack_steps.items():
+        for node_dict in serialized_attack_steps.values():
 
             # Recreate asset links if model is available.
             node_asset = None
@@ -306,8 +149,14 @@ class AttackGraph():
                     logger.error(msg, node_dict["asset"])
                     raise LookupError(msg % node_dict["asset"])
 
+            lg_asset_name, lg_attack_step_name = \
+                disaggregate_attack_step_full_name(
+                    node_dict['lang_graph_attack_step'])
+            lg_attack_step = lang_graph.assets[lg_asset_name].\
+                attack_steps[lg_attack_step_name]
             ag_node = AttackGraphNode(
                 type=node_dict['type'],
+                lang_graph_attack_step = lg_attack_step,
                 name=node_dict['name'],
                 ttc=node_dict['ttc'],
                 asset=node_asset
@@ -330,17 +179,15 @@ class AttackGraph():
                 'is_viable' in node_dict else True
             ag_node.is_necessary = node_dict['is_necessary'] == 'True' if \
                 'is_necessary' in node_dict else True
-            ag_node.mitre_info = str(node_dict['mitre_info']) if \
-                'mitre_info' in node_dict else None
-            ag_node.tags = node_dict['tags'] if \
-                'tags' in node_dict else []
+            ag_node.tags = set(node_dict['tags']) if \
+                'tags' in node_dict else set()
             ag_node.extras = node_dict.get('extras', {})
 
             # Add AttackGraphNode to AttackGraph
             attack_graph.add_node(ag_node, node_id=node_dict['id'])
 
         # Re-establish links between nodes.
-        for node_full_name, node_dict in serialized_attack_steps.items():
+        for node_dict in serialized_attack_steps.values():
             _ag_node = attack_graph.get_node_by_id(node_dict['id'])
             if not isinstance(_ag_node, AttackGraphNode):
                 msg = ('Failed to find node with id %s when loading'
@@ -366,7 +213,7 @@ class AttackGraph():
                         raise LookupError(msg % parent_id)
                     _ag_node.parents.append(parent)
 
-        for attacker_name, attacker in serialized_attackers.items():
+        for attacker in serialized_attackers.values():
             ag_attacker = Attacker(
                 name = attacker['name'],
                 entry_points = [],
@@ -388,7 +235,8 @@ class AttackGraph():
     def load_from_file(
             cls,
             filename: str,
-            model: Optional[Model]=None
+            lang_graph: LanguageGraph,
+            model: Optional[Model] = None
         ) -> AttackGraph:
         """Create from json or yaml file depending on file extension"""
         if model is not None:
@@ -404,7 +252,8 @@ class AttackGraph():
             serialized_attack_graph = load_dict_from_json_file(filename)
         else:
             raise ValueError('Unknown file extension, expected json/yml/yaml')
-        return cls._from_dict(serialized_attack_graph, model=model)
+        return cls._from_dict(serialized_attack_graph,
+            lang_graph, model = model)
 
     def get_node_by_id(self, node_id: int) -> Optional[AttackGraphNode]:
         """
@@ -432,7 +281,7 @@ class AttackGraph():
         The attack step node that matches the given full name.
         """
 
-        logger.debug(f'Looking up node with full name "{full_name}"')
+        logger.debug(f'Looking up node with full name "%s"', full_name)
         return self._full_name_to_node.get(full_name)
 
     def get_attacker_by_id(self, attacker_id: int) -> Optional[Attacker]:
@@ -493,6 +342,170 @@ class AttackGraph():
 
             attacker.entry_points = list(attacker.reached_attack_steps)
 
+    def _follow_expr_chain(
+            self,
+            model: Model,
+            target_assets: set[Any],
+            expr_chain: Optional[ExpressionsChain]
+        ) -> set[Any]:
+        """
+        Recursively follow a language graph expressions chain on an instance
+        model.
+
+        Arguments:
+        model           - a maltoolbox.model.Model on which to follow the
+                          expressions chain
+        target_assets   - the set of assets that this expressions chain
+                          should apply to. Initially it will contain the
+                          asset to which the attack step belongs
+        expr_chain      - the expressions chain we are following
+
+        Return:
+        A list of all of the target assets.
+        """
+
+        if expr_chain is None:
+            # There is no expressions chain link left to follow return the
+            # current target assets
+            return set(target_assets)
+
+        if logger.isEnabledFor(logging.DEBUG):
+            # Avoid running json.dumps when not in debug
+            logger.debug(
+                'Following Expressions Chain:\n%s',
+                json.dumps(expr_chain.to_dict(), indent = 2)
+            )
+
+        match (expr_chain.type):
+            case 'union' | 'intersection' | 'difference':
+                # The set operators are used to combine the left hand and
+                # right hand targets accordingly.
+                if not expr_chain.left_link:
+                    raise LanguageGraphException('"%s" step expression chain'
+                        ' is missing the left link.' % expr_chain.type)
+                if not expr_chain.right_link:
+                    raise LanguageGraphException('"%s" step expression chain'
+                        ' is missing the right link.' % expr_chain.type)
+                lh_targets = self._follow_expr_chain(
+                    model,
+                    target_assets,
+                    expr_chain.left_link
+                )
+                rh_targets = self._follow_expr_chain(
+                    model,
+                    target_assets,
+                    expr_chain.right_link
+                )
+
+                match (expr_chain.type):
+                    # Once the assets become hashable set operations should be
+                    # used instead.
+                    case 'union':
+                        new_target_assets = lh_targets.union(rh_targets)
+
+                    case 'intersection':
+                        new_target_assets = lh_targets.intersection(rh_targets)
+
+                    case 'difference':
+                        new_target_assets = lh_targets.difference(rh_targets)
+
+                return new_target_assets
+
+            case 'field':
+                # Change the target assets from the current ones to the
+                # associated assets given the specified field name.
+                if not expr_chain.fieldname:
+                    raise LanguageGraphException('"field" step expression '
+                        'chain is missing fieldname.')
+                new_target_assets = set()
+                new_target_assets.update(
+                    *(
+                        model.get_associated_assets_by_field_name(
+                            asset, expr_chain.fieldname
+                        )
+                        for asset in target_assets
+                      )
+                )
+                return new_target_assets
+
+            case 'transitive':
+                if not expr_chain.sub_link:
+                    raise LanguageGraphException('"transitive" step '
+                        'expression chain is missing sub link.')
+
+                new_assets = target_assets
+
+                while new_assets := self._follow_expr_chain(
+                    model, new_assets, expr_chain.sub_link
+                ):
+                    if not (new_assets := new_assets.difference(target_assets)):
+                        break
+
+                    target_assets.update(new_assets)
+
+                return target_assets
+
+            case 'subType':
+                if not expr_chain.sub_link:
+                    raise LanguageGraphException('"subType" step '
+                        'expression chain is missing sub link.')
+                new_target_assets = set()
+                new_target_assets.update(
+                    self._follow_expr_chain(
+                        model, target_assets, expr_chain.sub_link
+                    )
+                )
+
+                selected_new_target_assets = set()
+                for asset in new_target_assets:
+                    lang_graph_asset = self.lang_graph.assets[asset.type]
+                    if not lang_graph_asset:
+                        raise LookupError(
+                            f'Failed to find asset \"{asset.type}\" in the '
+                            'language graph.'
+                        )
+                    lang_graph_subtype_asset = expr_chain.subtype
+                    if not lang_graph_subtype_asset:
+                        raise LookupError(
+                            'Failed to find asset "%s" in the '
+                            'language graph.' % expr_chain.subtype
+                        )
+                    if lang_graph_asset.is_subasset_of(
+                            lang_graph_subtype_asset):
+                        selected_new_target_assets.add(asset)
+
+                return selected_new_target_assets
+
+            case 'collect':
+                if not expr_chain.left_link:
+                    raise LanguageGraphException('"collect" step expression chain'
+                        ' is missing the left link.')
+                if not expr_chain.right_link:
+                    raise LanguageGraphException('"collect" step expression chain'
+                        ' is missing the right link.')
+                lh_targets = self._follow_expr_chain(
+                    model,
+                    target_assets,
+                    expr_chain.left_link
+                )
+                rh_targets = self._follow_expr_chain(
+                    model,
+                    lh_targets,
+                    expr_chain.right_link
+                )
+                return rh_targets
+
+            case _:
+                msg = 'Unknown attack expressions chain type: %s'
+                logger.error(
+                    msg,
+                    expr_chain.type
+                )
+                raise AttackGraphStepExpressionError(
+                    msg % expr_chain.type
+                )
+                return None
+
     def _generate_graph(self) -> None:
         """
         Generate the attack graph based on the original model instance and the
@@ -514,57 +527,67 @@ class AttackGraph():
 
             attack_step_nodes = []
 
-            # TODO probably part of what happens here is already done in lang_graph
-            attack_steps = self.lang_graph._get_attacks_for_asset_type(asset.type)
+            lang_graph_asset = self.lang_graph.assets[asset.type]
+            if lang_graph_asset is None:
+                raise LookupError(
+                    f'Failed to find asset with name \"{asset.type}\" in '
+                    'the language graph.'
+                )
 
-            for attack_step_name, attack_step_attribs in attack_steps.items():
+            for attack_step in lang_graph_asset.attack_steps.values():
                 logger.debug(
-                    'Generating attack step node for %s.', attack_step_name
+                    'Generating attack step node for %s.', attack_step.name
                 )
 
                 defense_status = None
                 existence_status = None
-                node_name = asset.name + ':' + attack_step_name
+                node_name = asset.name + ':' + attack_step.name
 
-                match (attack_step_attribs['type']):
+                match (attack_step.type):
                     case 'defense':
                         # Set the defense status for defenses
-                        defense_status = getattr(asset, attack_step_name)
+                        defense_status = getattr(asset, attack_step.name)
                         logger.debug(
-                            'Setting the defense status of %s to %s.',
+                            'Setting the defense status of \"%s\" to '
+                            '\"%s\".',
                             node_name, defense_status
                         )
 
                     case 'exist' | 'notExist':
-                        # Resolve step expression associated with (non-)existence
-                        # attack steps.
-                        (target_assets, attack_step) = _process_step_expression(
-                            self.lang_graph,
-                            self.model,
-                            [asset],
-                            attack_step_attribs['requires']['stepExpressions'][0])
-                        # If the step expression resolution yielded the target
-                        # assets then the required assets exist in the model.
-                        existence_status = target_assets != []
-
-                mitre_info = attack_step_attribs['meta']['mitre'] if 'mitre' in\
-                    attack_step_attribs['meta'] else None
+                        # Resolve step expression associated with
+                        # (non-)existence attack steps.
+                        existence_status = False
+                        for requirement in attack_step.requires:
+                            target_assets = self._follow_expr_chain(
+                                    self.model,
+                                    set([asset]),
+                                    requirement
+                                )
+                            # If the step expression resolution yielded
+                            # the target assets then the required assets
+                            # exist in the model.
+                            if target_assets:
+                                existence_status = True
+                                break
+
+                    case _:
+                        pass
+
                 ag_node = AttackGraphNode(
-                    type = attack_step_attribs['type'],
+                    type = attack_step.type,
+                    lang_graph_attack_step = attack_step,
                     asset = asset,
-                    name = attack_step_name,
-                    ttc = attack_step_attribs['ttc'],
+                    name = attack_step.name,
+                    ttc = attack_step.ttc,
                     children = [],
                     parents = [],
                     defense_status = defense_status,
                     existence_status = existence_status,
                     is_viable = True,
                     is_necessary = True,
-                    mitre_info = mitre_info,
-                    tags = attack_step_attribs['tags'],
+                    tags = set(attack_step.tags),
                     compromised_by = []
                 )
-                ag_node.attributes = attack_step_attribs
                 attack_step_nodes.append(ag_node)
                 self.add_node(ag_node)
             asset.attack_step_nodes = attack_step_nodes
@@ -576,42 +599,66 @@ class AttackGraph():
                 ag_node.full_name,
                 ag_node.id
             )
-            step_expressions = \
-                ag_node.attributes['reaches']['stepExpressions'] if \
-                    isinstance(ag_node.attributes, dict) and ag_node.attributes['reaches'] else []
-
-            for step_expression in step_expressions:
-                # Resolve each of the attack step expressions listed for this
-                # attack step to determine children.
-                (target_assets, attack_step) = _process_step_expression(
-                    self.lang_graph,
-                    self.model,
-                    [ag_node.asset],
-                    step_expression)
-
-                for target in target_assets:
-                    target_node_full_name = target.name + ':' + attack_step
-                    target_node = self.get_node_by_full_name(
-                        target_node_full_name
-                    )
-                    if not target_node:
-                        msg = ('Failed to find target node '
-                               '"%s" to link with for attack step "%s"(%d)!')
-                        logger.error(
-                            msg,
-                            target_node_full_name,
-                            ag_node.full_name,
-                            ag_node.id
-                        )
-                        raise AttackGraphStepExpressionError(
-                            msg % (
-                                target_node_full_name,
-                                ag_node.full_name,
-                                ag_node.id
-                            )
+
+            if not ag_node.asset:
+                raise AttackGraphException('Attack graph node is missing '
+                    'asset link')
+            lang_graph_asset = self.lang_graph.assets[ag_node.asset.type]
+
+            lang_graph_attack_step = lang_graph_asset.attack_steps[\
+                ag_node.name]
+
+            while lang_graph_attack_step:
+                for child in lang_graph_attack_step.children.values():
+                    for target_attack_step, expr_chain in child:
+                        target_assets = self._follow_expr_chain(
+                            self.model,
+                            set([ag_node.asset]),
+                            expr_chain
                         )
-                    ag_node.children.append(target_node)
-                    target_node.parents.append(ag_node)
+
+                        for target_asset in target_assets:
+                            if target_asset is not None:
+                                target_node_full_name = target_asset.name + \
+                                    ':' + target_attack_step.name
+                                target_node = self.get_node_by_full_name(
+                                    target_node_full_name)
+                                if target_node is None:
+                                    msg = ('Failed to find target node '
+                                           '"%s" to link with for attack '
+                                           'step "%s"(%d)!')
+                                    logger.error(
+                                        msg,
+                                        target_node_full_name,
+                                        ag_node.full_name,
+                                        ag_node.id
+                                    )
+                                    raise AttackGraphStepExpressionError(
+                                        msg % (
+                                            target_node_full_name,
+                                            ag_node.full_name,
+                                            ag_node.id
+                                        )
+                                    )
+
+                                assert ag_node.id is not None
+                                assert target_node.id is not None
+
+                                logger.debug('Linking attack step "%s"(%d) '
+                                    'to attack step "%s"(%d)' %
+                                    (
+                                        ag_node.full_name,
+                                        ag_node.id,
+                                        target_node.full_name,
+                                        target_node.id
+                                    )
+                                )
+                                ag_node.children.append(target_node)
+                                target_node.parents.append(ag_node)
+                if lang_graph_attack_step.overrides:
+                    break
+                lang_graph_attack_step = lang_graph_attack_step.inherits
+
 
     def regenerate_graph(self) -> None:
         """
@@ -687,16 +734,18 @@ class AttackGraph():
         reached_attack_steps    - list of ids of the attack steps that the
                                   attacker has reached
         """
+
         if logger.isEnabledFor(logging.DEBUG):
             # Avoid running json.dumps when not in debug
             if attacker_id is not None:
                 logger.debug('Add attacker "%s" with id:%d.',
                     attacker.name,
-                    attacker_id)
+                    attacker_id
+                )
             else:
                 logger.debug('Add attacker "%s" without id.',
-                    attacker.name)
-
+                    attacker.name
+                )
 
         attacker.id = attacker_id or self.next_attacker_id
         if attacker.id in self._id_to_attacker: