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

maltoolbox/attackgraph/attackgraph.py

@@ -5,17 +5,21 @@ from __future__ import annotations
 import copy
 import logging
 import json
+import sys
+import zipfile
 
 from itertools import chain
 from typing import TYPE_CHECKING
 
+from .analyzers.apriori import calculate_viability_and_necessity
 from .node import AttackGraphNode
 from .attacker import Attacker
+from .. import log_configs
 from ..exceptions import AttackGraphStepExpressionError, AttackGraphException
 from ..exceptions import LanguageGraphException
 from ..model import Model
 from ..language import (LanguageGraph, ExpressionsChain,
-    disaggregate_attack_step_full_name)
+    LanguageGraphAttackStep, disaggregate_attack_step_full_name)
 from ..file_utils import (
     load_dict_from_json_file,
     load_dict_from_yaml_file,
@@ -25,20 +29,64 @@ from ..file_utils import (
 
 if TYPE_CHECKING:
     from typing import Any, Optional
+    from ..model import ModelAsset
 
 logger = logging.getLogger(__name__)
 
 
+def create_attack_graph(
+        lang_file: str,
+        model_file: str,
+        attach_attackers=True,
+        calc_viability_and_necessity=True
+    ) -> AttackGraph:
+    """Create and return an attack graph
+
+    Args:
+    lang_file                       - path to language file (.mar or .mal)
+    model_file                      - path to model file (yaml or json)
+    attach_attackers                - whether to run attach_attackers or not
+    calc_viability_and_necessity    - whether run apriori calculations or not
+    """
+    try:
+        lang_graph = LanguageGraph.from_mar_archive(lang_file)
+    except zipfile.BadZipFile:
+        lang_graph = LanguageGraph.from_mal_spec(lang_file)
+
+    if log_configs['langspec_file']:
+        lang_graph.save_to_file(log_configs['langspec_file'])
+
+    instance_model = Model.load_from_file(model_file, lang_graph)
+
+    if log_configs['model_file']:
+        instance_model.save_to_file(log_configs['model_file'])
+
+    try:
+        attack_graph = AttackGraph(lang_graph, instance_model)
+    except AttackGraphStepExpressionError:
+        logger.error(
+            'Attack graph generation failed when attempting '
+            'to resolve attack step expression!'
+        )
+        sys.exit(1)
+
+    if attach_attackers:
+        attack_graph.attach_attackers()
+
+    if calc_viability_and_necessity:
+        calculate_viability_and_necessity(attack_graph)
+
+    return attack_graph
+
+
 class AttackGraph():
     """Graph representation of attack steps"""
     def __init__(self, lang_graph, model: Optional[Model] = None):
-        self.nodes: list[AttackGraphNode] = []
-        self.attackers: list[Attacker] = []
+        self.nodes: dict[int, AttackGraphNode] = {}
+        self.attackers: dict[int, Attacker] = {}
         # Dictionaries used in optimization to get nodes and attackers by id
         # or full name faster
-        self._id_to_node: dict[int, AttackGraphNode] = {}
         self._full_name_to_node: dict[str, AttackGraphNode] = {}
-        self._id_to_attacker: dict[int, Attacker] = {}
 
         self.model = model
         self.lang_graph = lang_graph
@@ -48,16 +96,17 @@ class AttackGraph():
             self._generate_graph()
 
     def __repr__(self) -> str:
-        return f'AttackGraph({len(self.nodes)} nodes)'
+        return (f'AttackGraph(Number of nodes: {len(self.nodes)}, '
+            f'model: {self.model}, language: {self.lang_graph}')
 
     def _to_dict(self) -> dict:
         """Convert AttackGraph to dict"""
         serialized_attack_steps = {}
         serialized_attackers = {}
-        for ag_node in self.nodes:
+        for ag_node in self.nodes.values():
             serialized_attack_steps[ag_node.full_name] =\
                 ag_node.to_dict()
-        for attacker in self.attackers:
+        for attacker in self.attackers.values():
             serialized_attackers[attacker.name] = attacker.to_dict()
         logger.debug('Serialized %d attack steps and %d attackers.' %
             (len(self.nodes), len(self.attackers))
@@ -76,34 +125,32 @@ class AttackGraph():
         copied_attackgraph = AttackGraph(self.lang_graph)
         copied_attackgraph.model = self.model
 
-        copied_attackgraph.nodes = []
+        copied_attackgraph.nodes = {}
 
         # Deep copy nodes
-        for node in self.nodes:
+        for node_id, node in self.nodes.items():
             copied_node = copy.deepcopy(node, memo)
-            copied_attackgraph.nodes.append(copied_node)
+            copied_attackgraph.nodes[node_id] = copied_node
 
         # Re-link node references
-        for node in self.nodes:
+        for node in self.nodes.values():
             if node.parents:
                 memo[id(node)].parents = copy.deepcopy(node.parents, memo)
             if node.children:
                 memo[id(node)].children = copy.deepcopy(node.children, memo)
 
-        # Deep copy attackers and references to them
-        copied_attackgraph.attackers = copy.deepcopy(self.attackers, memo)
+        # Deep copy attackers
+        for attacker_id, attacker in self.attackers.items():
+            copied_attacker = copy.deepcopy(attacker, memo)
+            copied_attackgraph.attackers[attacker_id] = copied_attacker
 
         # Re-link attacker references
-        for node in self.nodes:
+        for node in self.nodes.values():
             if node.compromised_by:
                 memo[id(node)].compromised_by = copy.deepcopy(
                     node.compromised_by, memo)
 
         # Copy lookup dicts
-        copied_attackgraph._id_to_attacker = \
-            copy.deepcopy(self._id_to_attacker, memo)
-        copied_attackgraph._id_to_node = \
-            copy.deepcopy(self._id_to_node, memo)
         copied_attackgraph._full_name_to_node = \
             copy.deepcopy(self._full_name_to_node, memo)
 
@@ -144,8 +191,8 @@ class AttackGraph():
             if model and 'asset' in node_dict:
                 node_asset = model.get_asset_by_name(node_dict['asset'])
                 if node_asset is None:
-                    msg = ('Failed to find asset with id %s'
-                            'when loading from attack graph dict')
+                    msg = ('Failed to find asset with name "%s"'
+                            ' when loading from attack graph dict')
                     logger.error(msg, node_dict["asset"])
                     raise LookupError(msg % node_dict["asset"])
 
@@ -154,13 +201,15 @@ class AttackGraph():
                     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
+            ag_node = attack_graph.add_node(
+                lg_attack_step = lg_attack_step,
+                node_id = node_dict['id'],
+                model_asset = node_asset,
+                defense_status = node_dict.get('defense_status', None),
+                existence_status = node_dict.get('existence_status', None)
             )
+            ag_node.tags = set(node_dict.get('tags', []))
+            ag_node.extras = node_dict.get('extras', {})
 
             if node_asset:
                 # Add AttackGraphNode to attack_step_nodes of asset
@@ -171,24 +220,10 @@ class AttackGraph():
                 else:
                     node_asset.attack_step_nodes = [ag_node]
 
-            ag_node.defense_status = float(node_dict['defense_status']) if \
-                'defense_status' in node_dict else None
-            ag_node.existence_status = node_dict['existence_status'] \
-                == 'True' if 'existence_status' in node_dict else None
-            ag_node.is_viable = node_dict['is_viable'] == 'True' if \
-                '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.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_dict in serialized_attack_steps.values():
-            _ag_node = attack_graph.get_node_by_id(node_dict['id'])
+            _ag_node = attack_graph.nodes[node_dict['id']]
             if not isinstance(_ag_node, AttackGraphNode):
                 msg = ('Failed to find node with id %s when loading'
                        ' attack graph from dict')
@@ -196,33 +231,36 @@ class AttackGraph():
                 raise LookupError(msg % node_dict["id"])
             else:
                 for child_id in node_dict['children']:
-                    child = attack_graph.get_node_by_id(int(child_id))
+                    child = attack_graph.nodes[int(child_id)]
                     if child is None:
                         msg = ('Failed to find child node with id %s'
                                ' when loading from attack graph from dict')
                         logger.error(msg, child_id)
                         raise LookupError(msg % child_id)
-                    _ag_node.children.append(child)
+                    _ag_node.children.add(child)
 
                 for parent_id in node_dict['parents']:
-                    parent = attack_graph.get_node_by_id(int(parent_id))
+                    parent = attack_graph.nodes[int(parent_id)]
                     if parent is None:
                         msg = ('Failed to find parent node with id %s '
                                'when loading from attack graph from dict')
                         logger.error(msg, parent_id)
                         raise LookupError(msg % parent_id)
-                    _ag_node.parents.append(parent)
+                    _ag_node.parents.add(parent)
 
         for attacker in serialized_attackers.values():
             ag_attacker = Attacker(
                 name = attacker['name'],
-                entry_points = [],
-                reached_attack_steps = []
+                entry_points = set(),
+                reached_attack_steps = set()
             )
             attack_graph.add_attacker(
                 attacker = ag_attacker,
                 attacker_id = int(attacker['id']),
-                entry_points = attacker['entry_points'].keys(),
+                entry_points = [
+                    int(node_id) # Convert to int since they can be strings
+                    for node_id in attacker['entry_points'].keys()
+                ],
                 reached_attack_steps = [
                     int(node_id) # Convert to int since they can be strings
                     for node_id in attacker['reached_attack_steps'].keys()
@@ -255,20 +293,6 @@ class AttackGraph():
         return cls._from_dict(serialized_attack_graph,
             lang_graph, model = model)
 
-    def get_node_by_id(self, node_id: int) -> Optional[AttackGraphNode]:
-        """
-        Return the attack node that matches the id provided.
-
-        Arguments:
-        node_id     - the id of the attack graph node we are looking for
-
-        Return:
-        The attack step node that matches the given id.
-        """
-
-        logger.debug('Looking up node with id %s', node_id)
-        return self._id_to_node.get(node_id)
-
     def get_node_by_full_name(self, full_name: str) -> Optional[AttackGraphNode]:
         """
         Return the attack node that matches the full name provided.
@@ -284,20 +308,6 @@ class AttackGraph():
         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]:
-        """
-        Return the attacker that matches the id provided.
-
-        Arguments:
-        attacker_id     - the id of the attacker we are looking for
-
-        Return:
-        The attacker that matches the given id.
-        """
-
-        logger.debug(f'Looking up attacker with id {attacker_id}')
-        return self._id_to_attacker.get(attacker_id)
-
     def attach_attackers(self) -> None:
         """
         Create attackers and their entry point nodes and attach them to the
@@ -322,8 +332,8 @@ class AttackGraph():
 
             attacker = Attacker(
                 name = attacker_info.name,
-                entry_points = [],
-                reached_attack_steps = []
+                entry_points = set(),
+                reached_attack_steps = set()
             )
             self.add_attacker(attacker)
 
@@ -340,12 +350,12 @@ class AttackGraph():
                         continue
                     attacker.compromise(ag_node)
 
-            attacker.entry_points = list(attacker.reached_attack_steps)
+            attacker.entry_points = set(attacker.reached_attack_steps)
 
     def _follow_expr_chain(
             self,
             model: Model,
-            target_assets: set[Any],
+            target_assets: set[ModelAsset],
             expr_chain: Optional[ExpressionsChain]
         ) -> set[Any]:
         """
@@ -420,10 +430,9 @@ class AttackGraph():
                 new_target_assets = set()
                 new_target_assets.update(
                     *(
-                        model.get_associated_assets_by_field_name(
-                            asset, expr_chain.fieldname
-                        )
-                        for asset in target_assets
+                        asset.associated_assets.get(
+                            expr_chain.fieldname, set()
+                        ) for asset in target_assets
                       )
                 )
                 return new_target_assets
@@ -518,7 +527,7 @@ class AttackGraph():
             raise AttackGraphException(msg)
 
         # First, generate all of the nodes of the attack graph.
-        for asset in self.model.assets:
+        for asset in self.model.assets.values():
 
             logger.debug(
                 'Generating attack steps for asset %s which is of class %s.',
@@ -527,14 +536,7 @@ class AttackGraph():
 
             attack_step_nodes = []
 
-            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 in lang_graph_asset.attack_steps.values():
+            for attack_step in asset.lg_asset.attack_steps.values():
                 logger.debug(
                     'Generating attack step node for %s.', attack_step.name
                 )
@@ -546,10 +548,9 @@ class AttackGraph():
                 match (attack_step.type):
                     case 'defense':
                         # Set the defense status for defenses
-                        defense_status = getattr(asset, attack_step.name)
+                        defense_status = asset.defenses[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
                         )
 
@@ -570,42 +571,40 @@ class AttackGraph():
                                 existence_status = True
                                 break
 
+                        logger.debug(
+                            'Setting the existence status of \"%s\" to '
+                            '%s.',
+                            node_name, existence_status
+                        )
+
                     case _:
                         pass
 
-                ag_node = AttackGraphNode(
-                    type = attack_step.type,
-                    lang_graph_attack_step = attack_step,
-                    asset = asset,
-                    name = attack_step.name,
-                    ttc = attack_step.ttc,
-                    children = [],
-                    parents = [],
+                ag_node = self.add_node(
+                    lg_attack_step = attack_step,
+                    model_asset = asset,
                     defense_status = defense_status,
-                    existence_status = existence_status,
-                    is_viable = True,
-                    is_necessary = True,
-                    tags = set(attack_step.tags),
-                    compromised_by = []
+                    existence_status = existence_status
                 )
                 attack_step_nodes.append(ag_node)
-                self.add_node(ag_node)
+
             asset.attack_step_nodes = attack_step_nodes
 
         # Then, link all of the nodes according to their associations.
-        for ag_node in self.nodes:
+        for ag_node in self.nodes.values():
             logger.debug(
                 'Determining children for attack step "%s"(%d)',
                 ag_node.full_name,
                 ag_node.id
             )
 
-            if not ag_node.asset:
+            if not ag_node.model_asset:
                 raise AttackGraphException('Attack graph node is missing '
                     'asset link')
-            lang_graph_asset = self.lang_graph.assets[ag_node.asset.type]
+            lang_graph_asset = self.lang_graph.assets[
+                ag_node.model_asset.type]
 
-            lang_graph_attack_step = lang_graph_asset.attack_steps[\
+            lang_graph_attack_step = lang_graph_asset.attack_steps[
                 ag_node.name]
 
             while lang_graph_attack_step:
@@ -613,7 +612,7 @@ class AttackGraph():
                     for target_attack_step, expr_chain in child:
                         target_assets = self._follow_expr_chain(
                             self.model,
-                            set([ag_node.asset]),
+                            set([ag_node.model_asset]),
                             expr_chain
                         )
 
@@ -653,8 +652,8 @@ class AttackGraph():
                                         target_node.id
                                     )
                                 )
-                                ag_node.children.append(target_node)
-                                target_node.parents.append(ag_node)
+                                ag_node.children.add(target_node)
+                                target_node.parents.add(ag_node)
                 if lang_graph_attack_step.overrides:
                     break
                 lang_graph_attack_step = lang_graph_attack_step.inherits
@@ -666,37 +665,68 @@ class AttackGraph():
         the MAL language specification provided at initialization.
         """
 
-        self.nodes = []
-        self.attackers = []
+        self.nodes = {}
+        self.attackers = {}
         self._generate_graph()
 
     def add_node(
             self,
-            node: AttackGraphNode,
-            node_id: Optional[int] = None
-        ) -> None:
-        """Add a node to the graph
+            lg_attack_step: LanguageGraphAttackStep,
+            node_id: Optional[int] = None,
+            model_asset: Optional[ModelAsset] = None,
+            defense_status: Optional[float] = None,
+            existence_status: Optional[bool] = None
+        ) -> AttackGraphNode:
+        """Create and add a node to the graph
         Arguments:
-        node    - the node to add
-        node_id - the id to assign to this node, usually used when loading
-                  an attack graph from a file
-        """
-        if logger.isEnabledFor(logging.DEBUG):
-            # Avoid running json.dumps when not in debug
-            logger.debug(f'Add node \"{node.full_name}\" '
-                f'with id:{node_id}:\n' \
-                + json.dumps(node.to_dict(), indent = 2))
+        lg_attack_step      - the language graph attack step that corresponds
+                              to the attack graph node to create
+        node_id             - id to assign to the newly created node, usually
+                              provided only when loading an existing attack
+                              graph from a file. If not provided the id will
+                              be set to the next highest id available.
+        model_asset         - the model asset that corresponds to the attack
+                              step node. While optional it is highly
+                              recommended that this be provided. It should
+                              only be ommitted if the model which was used to
+                              generate the attack graph is not available when
+                              loading an attack graph from a file.
+        defese_status       - the defense status of the node. Only, relevant
+                              for defense type nodes. A value between 0.0 and
+                              1.0 is expected.
+        existence_status    - the existence status of the node. Only, relevant
+                              for exist and notExist type nodes.
 
-        if node.id in self._id_to_node:
+        Return:
+        The newly created attack step node.
+        """
+        node_id = node_id if node_id is not None else self.next_node_id
+        if node_id in self.nodes:
             raise ValueError(f'Node index {node_id} already in use.')
+        self.next_node_id = max(node_id + 1, self.next_node_id)
 
-        node.id = node_id if node_id is not None else self.next_node_id
-        self.next_node_id = max(node.id + 1, self.next_node_id)
+        if logger.isEnabledFor(logging.DEBUG):
+            # Avoid running json.dumps when not in debug
+            logger.debug('Create and add to attackgraph node of type "%s" '
+                'with id:%d.\n' % (
+                    lg_attack_step.full_name,
+                    node_id
+                ))
+
+
+        node = AttackGraphNode(
+            node_id = node_id,
+            lg_attack_step = lg_attack_step,
+            model_asset = model_asset,
+            defense_status = defense_status,
+            existence_status = existence_status
+        )
 
-        self.nodes.append(node)
-        self._id_to_node[node.id] = node
+        self.nodes[node_id] = node
         self._full_name_to_node[node.full_name] = node
 
+        return node
+
     def remove_node(self, node: AttackGraphNode) -> None:
         """Remove node from attack graph
         Arguments:
@@ -709,11 +739,10 @@ class AttackGraph():
             child.parents.remove(node)
         for parent in node.parents:
             parent.children.remove(node)
-        self.nodes.remove(node)
 
         if not isinstance(node.id, int):
             raise ValueError(f'Invalid node id.')
-        del self._id_to_node[node.id]
+        del self.nodes[node.id]
         del self._full_name_to_node[node.full_name]
 
     def add_attacker(
@@ -748,12 +777,12 @@ class AttackGraph():
                 )
 
         attacker.id = attacker_id or self.next_attacker_id
-        if attacker.id in self._id_to_attacker:
+        if attacker.id in self.attackers:
             raise ValueError(f'Attacker index {attacker_id} already in use.')
 
         self.next_attacker_id = max(attacker.id + 1, self.next_attacker_id)
         for node_id in reached_attack_steps:
-            node = self.get_node_by_id(node_id)
+            node = self.nodes[node_id]
             if node:
                 attacker.compromise(node)
             else:
@@ -762,16 +791,15 @@ class AttackGraph():
                 logger.error(msg, node_id)
                 raise AttackGraphException(msg % node_id)
         for node_id in entry_points:
-            node = self.get_node_by_id(int(node_id))
+            node = self.nodes[node_id]
             if node:
-                attacker.entry_points.append(node)
+                attacker.entry_points.add(node)
             else:
                 msg = ("Could not find node with id %d"
                        "in attacker entrypoints.")
                 logger.error(msg, node_id)
                 raise AttackGraphException(msg % node_id)
-        self.attackers.append(attacker)
-        self._id_to_attacker[attacker.id] = attacker
+        self.attackers[attacker.id] = attacker
 
     def remove_attacker(self, attacker: Attacker):
         """Remove attacker from attack graph
@@ -785,7 +813,6 @@ class AttackGraph():
                 attacker.id)
         for node in attacker.reached_attack_steps:
             attacker.undo_compromise(node)
-        self.attackers.remove(attacker)
         if not isinstance(attacker.id, int):
             raise ValueError(f'Invalid attacker id.')
-        del self._id_to_attacker[attacker.id]
+        del self.attackers[attacker.id]