mal-toolbox 0.0.27__py3-none-any.whl → 0.1.12__py3-none-any.whl

maltoolbox/attackgraph/attackgraph.py

@@ -1,26 +1,44 @@
 """
 MAL-Toolbox Attack Graph Module
 """
-
+from __future__ import annotations
+import copy
 import logging
 import json
 
-from typing import List, Optional
+from typing import TYPE_CHECKING
+
+from .node import AttackGraphNode
+from .attacker import Attacker
+from ..exceptions import AttackGraphStepExpressionError
+from ..model import Model
+from ..exceptions import AttackGraphException
+from ..file_utils import (
+    load_dict_from_json_file,
+    load_dict_from_yaml_file,
+    save_dict_to_file
+)
 
-from maltoolbox.language import specification
-from maltoolbox.model import model
-from maltoolbox.attackgraph import node
-from maltoolbox.attackgraph import attacker
+if TYPE_CHECKING:
+    from typing import Any, Optional
+    from ..language import LanguageGraph
 
 logger = logging.getLogger(__name__)
 
-def _process_step_expression(lang: dict, model: model.Model,
-    target_assets: List, step_expression: dict):
+# 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            - a dictionary representing the MAL language specification
+    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
@@ -32,8 +50,13 @@ def _process_step_expression(lang: dict, model: model.Model,
     A tuple pair containing a list of all of the target assets and the name of
     the attack step.
     """
-    logger.debug('Processing Step Expression:\n' \
-        + json.dumps(step_expression, indent = 2))
+
+    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':
@@ -45,9 +68,9 @@ def _process_step_expression(lang: dict, model: model.Model,
             # The set operators are used to combine the left hand and right
             # hand targets accordingly.
             lh_targets, lh_attack_steps = _process_step_expression(
-                lang, model, target_assets, step_expression['lhs'])
+                lang_graph, model, target_assets, step_expression['lhs'])
             rh_targets, rh_attack_steps = _process_step_expression(
-                lang, model, target_assets, step_expression['rhs'])
+                lang_graph, model, target_assets, step_expression['rhs'])
 
             new_target_assets = []
             match (step_expression['type']):
@@ -60,7 +83,7 @@ def _process_step_expression(lang: dict, model: model.Model,
 
                 case 'intersection':
                     for ag_node in rh_targets:
-                        if next((lnode for lnode in new_target_assets \
+                        if next((lnode for lnode in lh_targets \
                             if lnode.id == ag_node.id), None):
                             new_target_assets.append(ag_node)
 
@@ -77,17 +100,19 @@ def _process_step_expression(lang: dict, model: model.Model,
             # 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, 'metaconcept')):
-                    variable_step_expr = specification.\
-                        get_variable_for_class_by_name(lang,
-                        target_asset.metaconcept, step_expression['name'])
+                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, model, target_assets, variable_step_expr)
+                        lang_graph, model, target_assets, variable_step_expr)
 
                 else:
-                    logger.error('Requested variable from non-asset'
-                        f'target node: {target_asset} which cannot be'
-                        'resolved.')
+                    logger.error(
+                        'Requested variable from non-asset target node:'
+                        '%s which cannot be resolved.', target_asset
+                    )
             return ([], None)
 
         case 'field':
@@ -112,7 +137,7 @@ def _process_step_expression(lang: dict, model: model.Model,
                         step_expression['stepExpression']['name']))
             if new_target_assets:
                 (additional_assets, _) = _process_step_expression(
-                    lang, model, new_target_assets, step_expression)
+                    lang_graph, model, new_target_assets, step_expression)
                 new_target_assets.extend(additional_assets)
                 return (new_target_assets, None)
             else:
@@ -122,91 +147,181 @@ def _process_step_expression(lang: dict, model: model.Model,
             new_target_assets = []
             for target_asset in target_assets:
                 (assets, _) = _process_step_expression(
-                    lang, model, target_assets, step_expression['stepExpression'])
+                    lang_graph, model, target_assets,
+                    step_expression['stepExpression'])
                 new_target_assets.extend(assets)
 
-            selected_new_target_assets = (asset for asset in \
-                new_target_assets if specification.extends_asset(
-                    lang,
-                    asset.metaconcept,
-                    step_expression['subType']))
+            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, model, target_assets, step_expression['lhs'])
-            return _process_step_expression(lang, model, lh_targets,
+                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: '
-                f'{step_expression["type"]}')
+            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):
+        self.nodes: list[AttackGraphNode] = []
+        self.attackers: list[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] = {}
 
-
-class AttackGraph:
-    def __init__(self, lang_spec = None, model: Optional[model.Model] = None):
-        self.nodes = []
-        self.attackers = []
         self.model = model
-        self.lang_spec = lang_spec
-        if self.model is not None and self.lang_spec is not None:
-            self.generate_graph(self.lang_spec, self.model)
-
-
-    def save_to_file(self, filename: str):
-        """
-        Save the attack graph to a json file.
-
-        Arguments:
-        filename        - the name of the output file
-        """
-
-        logger.info(f'Saving attack graph with {len(self.nodes)} attack step '
-            f'nodes to {filename} file.')
-        serialized_graph = []
+        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:
+            self._generate_graph()
+
+    def __repr__(self) -> str:
+        return f'AttackGraph({len(self.nodes)} nodes)'
+
+    def _to_dict(self) -> dict:
+        """Convert AttackGraph to dict"""
+        serialized_attack_steps = {}
+        serialized_attackers = {}
         for ag_node in self.nodes:
-            serialized_graph.append(ag_node.to_dict())
-        with open(filename, 'w', encoding='utf-8') as file:
-            json.dump(serialized_graph, file, indent=4)
-
-
-    def load_from_file(self, filename: str, model: Optional[model.Model] = None):
-        """
-        Load the attack graph model from a json file.
-
-        Arguments:
-        filename        - the name of the input file to parse
-        model           - (optional) the instance model that the attack graph was
-                          generated from. If this given then the attack graph node
-                          and instance model asset link can be re-established. If
-                          this argument is not given the attack graph will still
-                          be created it will just omit the links to the assets.
+            serialized_attack_steps[ag_node.full_name] =\
+                ag_node.to_dict()
+        for attacker in self.attackers:
+            serialized_attackers[attacker.name] = attacker.to_dict()
+        return {
+            'attack_steps': serialized_attack_steps,
+            'attackers': serialized_attackers,
+        }
+
+    def __deepcopy__(self, memo):
+
+        # Check if the object is already in the memo dictionary
+        if id(self) in memo:
+            return memo[id(self)]
+
+        copied_attackgraph = AttackGraph(self.lang_graph)
+        copied_attackgraph.model = self.model
+
+        copied_attackgraph.nodes = []
+
+        # Deep copy nodes
+        for node in self.nodes:
+            copied_node = copy.deepcopy(node, memo)
+            copied_attackgraph.nodes.append(copied_node)
+
+        # Re-link node references
+        for node in self.nodes:
+            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)
+
+        # Re-link attacker references
+        for node in self.nodes:
+            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)
+
+        # Copy counters
+        copied_attackgraph.next_node_id = self.next_node_id
+        copied_attackgraph.next_attacker_id = self.next_attacker_id
+
+        return copied_attackgraph
+
+    def save_to_file(self, filename: str) -> None:
+        """Save to json/yml depending on extension"""
+        logger.debug('Save attack graph to file "%s".', filename)
+        return save_dict_to_file(filename, self._to_dict())
+
+    @classmethod
+    def _from_dict(
+            cls,
+            serialized_object: dict,
+            model: Optional[Model]=None
+        ) -> AttackGraph:
+        """Create AttackGraph from dict
+        Args:
+        serialized_object   - AttackGraph in dict format
+        model               - Optional Model to add connections to
         """
 
-        logger.info(f'Loading attack graph from {filename} file.')
-        if model:
-            logger.info(f'Model(\'{model.name}\') was provided will attempt '
-            'to establish links to assets.')
-        else:
-            logger.info('No model was provided therefore asset links will '
-            'not be established.')
+        attack_graph = AttackGraph()
+        attack_graph.model = model
+        serialized_attack_steps = serialized_object['attack_steps']
+        serialized_attackers = serialized_object['attackers']
 
-        with open(filename, 'r', encoding='utf-8') as file:
-            serialized_graph = json.load(file)
         # Create all of the nodes in the imported attack graph.
-        for node_dict in serialized_graph:
-            ag_node = node.AttackGraphNode(
-                id=node_dict['id'],
+        for node_full_name, node_dict in serialized_attack_steps.items():
+
+            # Recreate asset links if model is available.
+            node_asset = None
+            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')
+                    logger.error(msg, node_dict["asset"])
+                    raise LookupError(msg % node_dict["asset"])
+
+            ag_node = AttackGraphNode(
                 type=node_dict['type'],
                 name=node_dict['name'],
-                ttc=node_dict['ttc']
+                ttc=node_dict['ttc'],
+                asset=node_asset
             )
 
+            if node_asset:
+                # Add AttackGraphNode to attack_step_nodes of asset
+                if hasattr(node_asset, 'attack_step_nodes'):
+                    node_attack_steps = list(node_asset.attack_step_nodes)
+                    node_attack_steps.append(ag_node)
+                    node_asset.attack_step_nodes = node_attack_steps
+                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'] \
@@ -218,180 +333,214 @@ class AttackGraph:
             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 None
-            if ag_node.name == 'firstSteps':
-                # This is an attacker entry point node, recreate the attacker.
-                attacker_id = ag_node.id.split(':')[1]
-                ag_attacker = attacker.Attacker(
-                    id = str(attacker_id),
-                    entry_points = [],
-                    reached_attack_steps = [],
-                    node = ag_node
-                )
-                self.attackers.append(ag_attacker)
-                ag_node.attacker = ag_attacker
+                'tags' in node_dict else []
+            ag_node.extras = node_dict.get('extras', {})
 
-            self.nodes.append(ag_node)
+            # 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_graph:
-            _ag_node: Optional[node.AttackGraphNode] = self.get_node_by_id(node_dict['id'])
-            if not isinstance(_ag_node, node.AttackGraphNode):
-                logger.error(f'Failed to find node with id {node_dict["id"]}'
-                    f' when loading from attack graph from file {filename}')
+        for node_full_name, node_dict in serialized_attack_steps.items():
+            _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'
+                       ' attack graph from dict')
+                logger.error(msg, node_dict["id"])
+                raise LookupError(msg % node_dict["id"])
             else:
                 for child_id in node_dict['children']:
-                    child = self.get_node_by_id(child_id)
+                    child = attack_graph.get_node_by_id(int(child_id))
                     if child is None:
-                        logger.error(f'Failed to find child node with id {child_id}'
-                            f' when loading from attack graph from file {filename}')
-                        return 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)
 
-                    if isinstance(_ag_node.attacker, attacker.Attacker):
-                        # Relink the attacker related connections since the node
-                        # is an attacker entry point node.
-                        ag_attacker = _ag_node.attacker
-                        ag_attacker.entry_points.append(child)
-                        ag_attacker.compromise(child)
-
                 for parent_id in node_dict['parents']:
-                    parent = self.get_node_by_id(parent_id)
+                    parent = attack_graph.get_node_by_id(int(parent_id))
                     if parent is None:
-                        logger.error('Failed to find parent node with id '
-                            f'{parent_id} when loading from attack graph from '
-                            f'file {filename}')
-                        return 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)
 
-                # Also recreate asset links if model is available.
-                if model and 'asset' in node_dict:
-                    asset = model.get_asset_by_id(
-                        int(node_dict['asset'].split(':')[1]))
-                    if asset is None:
-                        logger.error('Failed to find asset with id '
-                            f'{node_dict["asset"]} when loading from attack graph '
-                            f'from file {filename}')
-                        return None
-                    _ag_node.asset = asset
-                    if hasattr(asset, 'attack_step_nodes'):
-                        attack_step_nodes = list(asset.attack_step_nodes)
-                        attack_step_nodes.append(_ag_node)
-                        asset.attack_step_nodes = attack_step_nodes
-                    else:
-                        asset.attack_step_nodes = [_ag_node]
-
-
-    def get_node_by_id(self, node_id: str) -> Optional[node.AttackGraphNode]:
+        for attacker_name, attacker in serialized_attackers.items():
+            ag_attacker = Attacker(
+                name = attacker['name'],
+                entry_points = [],
+                reached_attack_steps = []
+            )
+            attack_graph.add_attacker(
+                attacker = ag_attacker,
+                attacker_id = int(attacker['id']),
+                entry_points = 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()
+                ]
+            )
+
+        return attack_graph
+
+    @classmethod
+    def load_from_file(
+            cls,
+            filename: str,
+            model: Optional[Model]=None
+        ) -> AttackGraph:
+        """Create from json or yaml file depending on file extension"""
+        if model is not None:
+            logger.debug('Load attack graph from file "%s" with '
+            'model "%s".', filename, model.name)
+        else:
+            logger.debug('Load attack graph from file "%s" '
+            'without model.', filename)
+        serialized_attack_graph = None
+        if filename.endswith(('.yml', '.yaml')):
+            serialized_attack_graph = load_dict_from_yaml_file(filename)
+        elif filename.endswith('.json'):
+            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)
+
+    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 none we are looking for
+        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(f'Looking up node with id {node_id}')
-        return next((ag_node for ag_node in self.nodes \
-            if ag_node.id == node_id), None)
+        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.
 
+        Arguments:
+        full_name   - the full name of the attack graph node we are looking
+                      for
 
-    def attach_attackers(self, model: model.Model):
+        Return:
+        The attack step node that matches the given full name.
         """
-        Create attackers and their entry point nodes and attach them to the
-        relevant attack step nodes and to the attackers.
+
+        logger.debug(f'Looking up node with full name "{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:
-        model       - the instance model containing the attackers
+        attacker_id     - the id of the attacker we are looking for
+
+        Return:
+        The attacker that matches the given id.
         """
 
-        logger.info(f'Attach attackers from \'{model.name}\' model to the '
-            'graph.')
-        for attacker_info in model.attackers:
-            attacker_node = node.AttackGraphNode(
-                    id = 'Attacker:' + str(attacker_info.id) + ':firstSteps',
-                    type = 'or',
-                    asset = None,
-                    name = 'firstSteps',
-                    ttc = {},
-                    children = [],
-                    parents = [],
-                    compromised_by = []
-            )
+        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
+        relevant attack step nodes and to the attackers.
+        """
+
+        if not self.model:
+            msg = "Can not attach attackers without a model"
+            logger.error(msg)
+            raise AttackGraphException(msg)
+
+        logger.info(
+            'Attach attackers from "%s" model to the graph.', self.model.name
+        )
+
+        for attacker_info in self.model.attackers:
 
-            ag_attacker = attacker.Attacker(
-                id = str(attacker_info.id),
+            if not attacker_info.name:
+                msg = "Can not attach attacker without name"
+                logger.error(msg)
+                raise AttackGraphException(msg)
+
+            attacker = Attacker(
+                name = attacker_info.name,
                 entry_points = [],
-                reached_attack_steps = [],
-                node = attacker_node
+                reached_attack_steps = []
             )
-            attacker_node.attacker = ag_attacker
-            self.attackers.append(ag_attacker)
+            self.add_attacker(attacker)
 
             for (asset, attack_steps) in attacker_info.entry_points:
                 for attack_step in attack_steps:
-                    attack_step_id = asset.metaconcept + ':' \
-                        + str(asset.id) + ':' + attack_step
-                    ag_node = self.get_node_by_id(attack_step_id)
+                    full_name = asset.name + ':' + attack_step
+                    ag_node = self.get_node_by_full_name(full_name)
                     if not ag_node:
-                        logger.warning('Failed to find attacker entry point '
-                            + attack_step_id + ' for Attacker:'
-                            + ag_attacker.id + '.')
+                        logger.warning(
+                            'Failed to find attacker entry point '
+                            '%s for %s.',
+                            full_name, attacker.name
+                        )
                         continue
-                    ag_attacker.compromise(ag_node)
-
-            ag_attacker.entry_points = ag_attacker.reached_attack_steps
-            attacker_node.children = ag_attacker.entry_points
-            self.nodes.append(attacker_node)
+                    attacker.compromise(ag_node)
 
+            attacker.entry_points = list(attacker.reached_attack_steps)
 
-    def generate_graph(self, lang: Optional[dict] = None, model: Optional[model.Model] = None):
+    def _generate_graph(self) -> None:
         """
-        Generate attack graph starting from a model instance
-        and a MAL language specification
-
-        Arguments:
-        model           - a maltoolbox.model.Model instance
-        lang            - a dictionary representing the MAL language specification
+        Generate the attack graph based on the original model instance and the
+        MAL language specification provided at initialization.
         """
 
-        if model is not None:
-            self.model = model
-        if lang is not None:
-            self.lang_spec = lang
-        if self.model is None or self.lang_spec is None:
-            return
+        if not self.model:
+            msg = "Can not generate AttackGraph without model"
+            logger.error(msg)
+            raise AttackGraphException(msg)
 
         # First, generate all of the nodes of the attack graph.
         for asset in self.model.assets:
-            logger.debug(f'Generating attack steps for asset {asset.name} which '\
-                f'is of class {asset.metaconcept}.')
+
+            logger.debug(
+                'Generating attack steps for asset %s which is of class %s.',
+                asset.name, asset.type
+            )
+
             attack_step_nodes = []
-            attack_steps = specification.get_attacks_for_class(lang,
-                asset.metaconcept)
+
+            # 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)
+
             for attack_step_name, attack_step_attribs in attack_steps.items():
-                logger.debug('Generating attack step node for '\
-                    f'{attack_step_name}.')
+                logger.debug(
+                    'Generating attack step node for %s.', attack_step_name
+                )
 
                 defense_status = None
-                existence_status: Optional[bool] = None
-                node_id = asset.metaconcept + ':' + str(asset.id) + ':' + attack_step_name
+                existence_status = None
+                node_name = asset.name + ':' + attack_step_name
 
                 match (attack_step_attribs['type']):
                     case 'defense':
                         # Set the defense status for defenses
                         defense_status = getattr(asset, attack_step_name)
-                        logger.debug('Setting the defense status of '\
-                            f'{node_id} to {defense_status}.')
+                        logger.debug(
+                            '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(
-                            lang,
-                            model,
+                            self.lang_graph,
+                            self.model,
                             [asset],
                             attack_step_attribs['requires']['stepExpressions'][0])
                         # If the step expression resolution yielded the target
@@ -400,8 +549,7 @@ class AttackGraph:
 
                 mitre_info = attack_step_attribs['meta']['mitre'] if 'mitre' in\
                     attack_step_attribs['meta'] else None
-                ag_node = node.AttackGraphNode(
-                    id = node_id,
+                ag_node = AttackGraphNode(
                     type = attack_step_attribs['type'],
                     asset = asset,
                     name = attack_step_name,
@@ -418,13 +566,16 @@ class AttackGraph:
                 )
                 ag_node.attributes = attack_step_attribs
                 attack_step_nodes.append(ag_node)
-                self.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:
-            logger.debug('Determining children for attack step '\
-                f'{ag_node.id}.')
+            logger.debug(
+                'Determining children for attack step "%s"(%d)',
+                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 []
@@ -432,21 +583,160 @@ class AttackGraph:
             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(lang,
-                     model, [ag_node.asset], step_expression)
+                (target_assets, attack_step) = _process_step_expression(
+                    self.lang_graph,
+                    self.model,
+                    [ag_node.asset],
+                    step_expression)
+
                 for target in target_assets:
-                    target_node_id = target.metaconcept + ':' \
-                        + str(target.id) + ':' + attack_step
-                    target_node = self.get_node_by_id(target_node_id)
+                    target_node_full_name = target.name + ':' + attack_step
+                    target_node = self.get_node_by_full_name(
+                        target_node_full_name
+                    )
                     if not target_node:
-                        logger.error('Failed to find targed node ' \
-                        f'{target_node_id} to link with for attack step ' \
-                        f'{ag_node.id}!')
-                        print('Failed to find targed node ' \
-                        f'{target_node_id} to link with for attack step ' \
-                        f'{ag_node.id}!')
-                        return 1
+                        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
+                            )
+                        )
                     ag_node.children.append(target_node)
                     target_node.parents.append(ag_node)
 
-        return 0
+    def regenerate_graph(self) -> None:
+        """
+        Regenerate the attack graph based on the original model instance and
+        the MAL language specification provided at initialization.
+        """
+
+        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
+        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))
+
+        if node.id in self._id_to_node:
+            raise ValueError(f'Node index {node_id} already in use.')
+
+        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)
+
+        self.nodes.append(node)
+        self._id_to_node[node.id] = node
+        self._full_name_to_node[node.full_name] = node
+
+    def remove_node(self, node: AttackGraphNode) -> None:
+        """Remove node from attack graph
+        Arguments:
+        node    - the node we wish to remove from the attack graph
+        """
+        if logger.isEnabledFor(logging.DEBUG):
+            # Avoid running json.dumps when not in debug
+            logger.debug(f'Remove node "%s"(%d).', node.full_name, node.id)
+        for child in node.children:
+            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._full_name_to_node[node.full_name]
+
+    def add_attacker(
+            self,
+            attacker: Attacker,
+            attacker_id: Optional[int] = None,
+            entry_points: list[int] = [],
+            reached_attack_steps: list[int] = []
+        ):
+        """Add an attacker to the graph
+        Arguments:
+        attacker                - the attacker to add
+        attacker_id             - the id to assign to this attacker, usually
+                                  used when loading an attack graph from a
+                                  file
+        entry_points            - list of attack step ids that serve as entry
+                                  points for the attacker
+        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)
+            else:
+                logger.debug('Add attacker "%s" without id.',
+                    attacker.name)
+
+
+        attacker.id = attacker_id or self.next_attacker_id
+        if attacker.id in self._id_to_attacker:
+            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)
+            if node:
+                attacker.compromise(node)
+            else:
+                msg = ("Could not find node with id %d"
+                       "in reached attack steps.")
+                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))
+            if node:
+                attacker.entry_points.append(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
+
+    def remove_attacker(self, attacker: Attacker):
+        """Remove attacker from attack graph
+        Arguments:
+        attacker    - the attacker we wish to remove from the attack graph
+        """
+        if logger.isEnabledFor(logging.DEBUG):
+            # Avoid running json.dumps when not in debug
+            logger.debug('Remove attacker "%s" with id:%d.',
+                attacker.name,
+                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]