mal-toolbox 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

maltoolbox/model.py

@@ -1,41 +1,37 @@
-"""
-MAL-Toolbox Model Module
+"""MAL-Toolbox Model Module
 """
 
 from __future__ import annotations
+
 import json
 import logging
-from typing import TYPE_CHECKING
 import math
+from typing import TYPE_CHECKING
 
+from . import __version__
+from .exceptions import ModelException
 from .file_utils import (
     load_dict_from_json_file,
     load_dict_from_yaml_file,
-    save_dict_to_file
+    save_dict_to_file,
 )
-
-from . import __version__
-from .exceptions import ModelException
-
+from .language import LanguageGraph
 if TYPE_CHECKING:
-    from typing import Any, Optional
-    from .language import (
-        LanguageGraph,
-        LanguageGraphAsset,
-        LanguageGraphAssociation
-    )
+    from typing import Any
+
+    from .language import LanguageGraphAsset, LanguageGraphAssociation
 
 logger = logging.getLogger(__name__)
 
 
-class Model():
+class Model:
     """An implementation of a MAL language model containing assets"""
+
     next_id: int = 0
 
     def __repr__(self) -> str:
         return f'Model(name: "{self.name}", language: {self.lang_graph})'
 
-
     def __init__(
             self,
             name: str,
@@ -45,25 +41,24 @@ class Model():
 
         self.name = name
         self.assets: dict[int, ModelAsset] = {}
-        self._name_to_asset:dict[str, ModelAsset] = {} # optimization
+        self._name_to_asset: dict[str, ModelAsset] = {}  # optimization
         self.lang_graph = lang_graph
         self.maltoolbox_version: str = mt_version
 
-
     def add_asset(
             self,
             asset_type: str,
-            name: Optional[str] = None,
-            asset_id: Optional[int] = None,
-            defenses: Optional[dict[str, float]] = None,
-            extras: Optional[dict] = None,
+            name: str | None = None,
+            asset_id: int | None = None,
+            defenses: dict[str, float] | None = None,
+            extras: dict | None = None,
             allow_duplicate_names: bool = True
         ) -> ModelAsset:
-        """
-        Create an asset based on the provided parameters and add it to the
+        """Create an asset based on the provided parameters and add it to the
         model.
 
         Arguments:
+        ---------
         asset_type              - string containing the asset type name
         name                    - string containing the asset name. If not
                                   provided the concatenated asset type and id
@@ -79,9 +74,10 @@ class Model():
                                   be appended with the id.
 
         Return:
+        ------
         The newly created asset.
-        """
 
+        """
         # Set asset ID and check for duplicates
         asset_id = asset_id or self.next_id
         if asset_id in self.assets:
@@ -91,15 +87,14 @@ class Model():
 
         if not name:
             name = asset_type + ':' + str(asset_id)
-        else:
-            if name in self._name_to_asset:
-                if allow_duplicate_names:
-                    name = name + ':' + str(asset_id)
-                else:
-                    raise ValueError(
-                        f'Asset name {name} is a duplicate'
-                        ' and we do not allow duplicates.'
-                    )
+        elif name in self._name_to_asset:
+            if allow_duplicate_names:
+                name = name + ':' + str(asset_id)
+            else:
+                raise ValueError(
+                    f'Asset name {name} is a duplicate'
+                    ' and we do not allow duplicates.'
+                )
 
         if asset_type not in self.lang_graph.assets:
             raise ValueError(
@@ -110,11 +105,11 @@ class Model():
         lg_asset = self.lang_graph.assets[asset_type]
 
         asset = ModelAsset(
-            name = name,
-            asset_id = asset_id,
-            lg_asset = lg_asset,
-            defenses = defenses,
-            extras = extras)
+            name=name,
+            asset_id=asset_id,
+            lg_asset=lg_asset,
+            defenses=defenses,
+            extras=extras)
 
         logger.debug(
             'Add "%s"(%d) to model "%s".', name, asset_id, self.name
@@ -124,14 +119,14 @@ class Model():
 
         return asset
 
-
     def remove_asset(self, asset: ModelAsset) -> None:
         """Remove an asset from the model.
 
         Arguments:
+        ---------
         asset     - the asset to remove
-        """
 
+        """
         logger.debug(
             'Remove "%s"(%d) from model "%s".',
             asset.name, asset.id, self.name
@@ -152,18 +147,19 @@ class Model():
         del self.assets[asset.id]
         del self._name_to_asset[asset.name]
 
-
     def get_asset_by_id(
             self, asset_id: int
-        ) -> Optional[ModelAsset]:
-        """
-        Find an asset in the model based on its id.
+        ) -> ModelAsset | None:
+        """Find an asset in the model based on its id.
 
         Arguments:
+        ---------
         asset_id        - the id of the asset we are looking for
 
         Return:
+        ------
         An asset matching the id if it exists in the model.
+
         """
         logger.debug(
             'Get asset with id %d from model "%s".',
@@ -171,18 +167,19 @@ class Model():
         )
         return self.assets.get(asset_id, None)
 
-
     def get_asset_by_name(
             self, asset_name: str
-        ) -> Optional[ModelAsset]:
-        """
-        Find an asset in the model based on its name.
+        ) -> ModelAsset | None:
+        """Find an asset in the model based on its name.
 
         Arguments:
+        ---------
         asset_name        - the name of the asset we are looking for
 
         Return:
+        ------
         An asset matching the name if it exists in the model.
+
         """
         logger.debug(
             'Get asset with name "%s" from model "%s".',
@@ -216,13 +213,11 @@ class Model():
 
         return contents
 
-
     def save_to_file(self, filename: str) -> None:
         """Save to json/yml depending on extension"""
         logger.debug('Save instance model to file "%s".', filename)
         return save_dict_to_file(filename, self._to_dict())
 
-
     @classmethod
     def _from_dict(
         cls,
@@ -232,17 +227,18 @@ class Model():
         """Create a model from dict representation
 
         Arguments:
+        ---------
         serialized_object    - Model in dict format
         lang_graph -
-        """
 
+        """
         maltoolbox_version = serialized_object['metadata']['MAL Toolbox Version'] \
             if 'MAL Toolbox Version' in serialized_object['metadata'] \
             else __version__
         model = Model(
             serialized_object['metadata']['name'],
             lang_graph,
-            mt_version = maltoolbox_version)
+            mt_version=maltoolbox_version)
 
         # Reconstruct the assets
         for asset_id, asset_dict in serialized_object['assets'].items():
@@ -264,12 +260,12 @@ class Model():
             )
 
             model.add_asset(
-                asset_type = asset_dict['type'],
-                name = asset_dict['name'],
-                defenses = {defense: float(value) for defense, value in \
+                asset_type=asset_dict['type'],
+                name=asset_dict['name'],
+                defenses={defense: float(value) for defense, value in
                     asset_dict.get('defenses', {}).items()},
-                extras = asset_dict.get('extras', {}),
-                asset_id = int(asset_id))
+                extras=asset_dict.get('extras', {}),
+                asset_id=int(asset_id))
 
         # Reconstruct the association links
         for asset_id, asset_dict in serialized_object['assets'].items():
@@ -291,7 +287,6 @@ class Model():
 
         return model
 
-
     @classmethod
     def load_from_file(
         cls,
@@ -315,6 +310,27 @@ class Model():
                 "Try to upgrade it with 'maltoolbox upgrade-model'"
             ) from e
 
+    def __getstate__(self):
+        lang_state = self.lang_graph.__getstate__()
+        state = self._to_dict()
+        return {
+            'model_state': state,
+            'lang_graph': lang_state
+        }
+    
+    def __setstate__(self, state):
+        # Restore the language graph first
+        lang_graph = LanguageGraph.__new__(LanguageGraph)
+        lang_graph.__setstate__(state['lang_graph'])
+        self.lang_graph = lang_graph
+        
+        # Restore the model state by creating a temporary model and copying attributes
+        temp_model = self._from_dict(state['model_state'], self.lang_graph)
+        self.name = temp_model.name
+        self.assets = temp_model.assets
+        self._name_to_asset = temp_model._name_to_asset
+        self.maltoolbox_version = temp_model.maltoolbox_version
+        self.next_id = temp_model.next_id
 
 class ModelAsset:
     def __init__(
@@ -322,8 +338,8 @@ class ModelAsset:
         name: str,
         asset_id: int,
         lg_asset: LanguageGraphAsset,
-        defenses: Optional[dict[str, float]] = None,
-        extras: Optional[dict] = None
+        defenses: dict[str, float] | None = None,
+        extras: dict | None = None
     ):
 
         self.name: str = name
@@ -335,9 +351,8 @@ class ModelAsset:
         self._associated_assets: dict[str, set[ModelAsset]] = {}
         self.attack_step_nodes: list = []
 
-    def _to_dict(self):
+    def _to_dict(self) -> dict[int, dict[str, Any]]:
         """Get dictionary representation of the asset."""
-
         logger.debug(
             'Translating "%s"(%d) to dictionary.', self.name, self.id)
 
@@ -365,7 +380,6 @@ class ModelAsset:
 
         return {self.id: asset_dict}
 
-
     def __repr__(self):
         return (f'ModelAsset(name: "{self.name}", id: {self.id}, '
             f'type: {self.type})')
@@ -388,10 +402,8 @@ class ModelAsset:
         return assocs_in_common
 
     def has_association_with(self, b: ModelAsset, assoc_name: str) -> bool:
+        """Returns True if association `assoc_name` exists between self and `b`
         """
-        Returns True if association `assoc_name` exists between self and `b`
-        """
-
         for fieldname, associated_assets in self.associated_assets.items():
             assoc = self.lg_asset.associations[fieldname]
             if assoc.name == assoc_name and b in associated_assets:
@@ -402,20 +414,20 @@ class ModelAsset:
     def validate_associated_assets(
             self, fieldname: str, assets_to_add: set[ModelAsset]
         ):
-        """
-        Validate an association we want to add (through `fieldname`)
+        """Validate an association we want to add (through `fieldname`)
         is valid with the assets given in param `assets_to_add`:
         - fieldname is valid for the asset type of this ModelAsset
         - type of `assets_to_add` is valid for the association
         - no more assets than 'field.maximum' are added to the field
 
-        Raises:
+        Raises
+        ------
             LookupError - fieldname can not be found for this ModelAsset
             ValueError - there will be too many assets in the field
                          if we add this association
             TypeError - if the asset type of `assets_to_add` is not valid
-        """
 
+        """
         # Validate that the field name is allowed for this asset type
         if fieldname not in self.lg_asset.associations:
             accepted_fieldnames = list(self.lg_asset.associations.keys())
@@ -450,11 +462,9 @@ class ModelAsset:
             )
 
     def add_associated_assets(self, fieldname: str, assets: set[ModelAsset]):
-        """
-        Add the assets provided as a parameter to the set of associated
+        """Add the assets provided as a parameter to the set of associated
         assets dictionary entry corresponding to the given fieldname.
         """
-
         if fieldname not in self.lg_asset.associations:
             if assets:
                 to_asset_type = next(iter(assets)).lg_asset
@@ -490,7 +500,7 @@ class ModelAsset:
 
     def remove_associated_assets(
             self, fieldname: str, assets: set[ModelAsset]):
-        """ Remove the assets provided as a parameter from the set of
+        """Remove the assets provided as a parameter from the set of
         associated assets dictionary entry corresponding to the fieldname
         parameter.
         """
@@ -507,12 +517,10 @@ class ModelAsset:
         if len(self._associated_assets[fieldname]) == 0:
             del self._associated_assets[fieldname]
 
-
     @property
     def associated_assets(self):
         return self._associated_assets
 
-
     @property
     def id(self):
         return self._id

maltoolbox/patternfinder/attackgraph_patterns.py

@@ -1,13 +1,18 @@
 """Utilities for finding patterns in the AttackGraph"""
 
 from __future__ import annotations
+
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable
+
 from maltoolbox.attackgraph import AttackGraph, AttackGraphNode
 
+
 class SearchPattern:
     """A pattern consists of conditions, the conditions are used
-    to find all matching sequences of nodes in an AttackGraph."""
+    to find all matching sequences of nodes in an AttackGraph.
+    """
+
     conditions: list[SearchCondition]
 
     def __init__(self, conditions):
@@ -19,11 +24,12 @@ class SearchPattern:
         that match all the conditions in the pattern
 
         Args:
+        ----
         graph   - The AttackGraph to search in
     
         Return: list[list[AttackGraphNode]] matching paths of Nodes
-        """
 
+        """
         # Find the starting nodes which match the first condition
         condition = self.conditions[0]
         matching_paths = []
@@ -34,12 +40,14 @@ class SearchPattern:
                 )
         return matching_paths
 
-@dataclass
+
+@dataclass(frozen=True, eq=True)
 class SearchCondition:
     """A condition that has to be true for a node to match"""
 
     # Predefined search conditions
-    ANY = lambda _: True
+    @staticmethod
+    def ANY(_): return True
 
     # `matches` should be a lambda that takes node as input and returns bool
     # If lamdba returns True for a node, the node matches
@@ -64,7 +72,7 @@ def find_matches_recursively(
         node: AttackGraphNode,
         condition_list: list[SearchCondition],
         current_path: list[AttackGraphNode] | None = None,
-        matching_paths: set[tuple[AttackGraphNode,...]] | None = None,
+        matching_paths: set[tuple[AttackGraphNode, ...]] | None = None,
         condition_match_count: int = 0
     ):
     """Find all paths of nodes that match the list of conditions.
@@ -73,6 +81,7 @@ def find_matches_recursively(
     The function runs recursively down all paths of children nodes.
 
     Args:
+    ----
     node                  - node to check if current `condition` matches for
     condition_list        - first condition in list will attempt match `node`
     current_path          - list of matched nodes so far (recursively built)
@@ -80,8 +89,8 @@ def find_matches_recursively(
     condition_match_count - number of matches on current condition so far
 
     Return: set of tuples (paths) of AttackGraphNodes that match the condition
-    """
 
+    """
     # Init path lists if None, or copy/init into new lists for each iteration
     current_path = [] if current_path is None else list(current_path)
     matching_paths = set() if matching_paths is None else matching_paths
@@ -129,6 +138,6 @@ def find_matches_recursively(
 
         if not next_conds:
             # Congrats - matched a full unique search pattern!
-            matching_paths.add(tuple(current_path)) # tuple is hashable
+            matching_paths.add(tuple(current_path))  # tuple is hashable
 
     return matching_paths

maltoolbox/translators/__init__.py

@@ -0,0 +1,8 @@
+from .networkx import attack_graph_to_nx, model_to_nx
+from .updater import load_model_from_older_version
+
+__all__ = [
+    'attack_graph_to_nx',
+    'load_model_from_older_version',
+    'model_to_nx'
+]

maltoolbox/translators/networkx.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import networkx as nx
+
+from maltoolbox.attackgraph import AttackGraph, AttackGraphNode
+from maltoolbox.model import Model
+
+
+def attack_graph_to_nx(nodes: AttackGraph | Iterable[AttackGraphNode]) -> nx.DiGraph:
+    """Convert an attack graph to a networkx DiGraph"""
+    if isinstance(nodes, AttackGraph):
+        nodes = list(nodes.nodes.values())
+    G: nx.DiGraph = nx.DiGraph()
+
+    for node in nodes:
+        G.add_node(node.id, **node.to_dict())
+        G.nodes[node.id]["full_name"] = node.full_name
+
+    edges = [(node.id, child.id) for node in nodes for child in node.children]
+    edges += [(parent.id, node.id) for node in nodes for parent in node.parents]
+    G.add_edges_from(edges)
+
+    return G
+
+
+def model_to_nx(model: Model) -> nx.Graph:
+    """Convert a MAL Model to a networkx GRaph"""
+    G: nx.Graph = nx.Graph()
+
+    for id, asset in model.assets.items():
+        asset_dict = asset._to_dict()[id]
+        asset_dict["id"] = id
+        G.add_node(id, **asset_dict)
+
+    for id, asset in model.assets.items():
+        for fieldname, associated_assets in asset.associated_assets.items():
+            for associated_asset in associated_assets:
+                G.add_edge(id, associated_asset.id, name=asset.lg_asset.associations[fieldname].name)
+
+    return G

maltoolbox/translators/updater.py

@@ -1,25 +1,20 @@
-import json
 import logging
 
-import yaml
-
-import logging
-from ..model import Model
-from ..language import LanguageGraph
 from ..file_utils import load_dict_from_json_file, load_dict_from_yaml_file
+from ..language import LanguageGraph
+from ..model import Model
 
 logger = logging.getLogger(__name__)
 
+
 def load_model_from_older_version(
         filename: str, lang_graph: LanguageGraph,
     ) -> Model:
-
-    """ Load an older Model file
+    """Load an older Model file
 
     Load an older model from given `filename` (yml/json)
     convert the model to the new format and return a Model object. 
     """
-
     model_dict = load_model_dict_from_file(filename)
 
     # Get the version of the model, default to 0.0
@@ -54,7 +49,6 @@ def load_model_dict_from_file(
     filename: str,
     ) -> dict:
     """Load a json or yaml file to dict"""
-
     model_dict = {}
     if filename.endswith('.yml') or filename.endswith('.yaml'):
         model_dict = load_dict_from_yaml_file(filename)
@@ -68,16 +62,17 @@ def load_model_dict_from_file(
 
 
 def convert_model_dict_from_version_0_0(model_dict: dict) -> dict:
-    """
-    Convert model dict version 0.0 to 0.1
+    """Convert model dict version 0.0 to 0.1
 
     Arguments:
+    ---------
     model_dict  - the dictionary containing the serialized model
 
     Returns:
+    -------
     A dictionary containing the version 0.1 equivalent serialized model
-    """
 
+    """
     new_model_dict = {}
 
     # Meta data and attackers did not change
@@ -121,24 +116,21 @@ def convert_model_dict_from_version_0_0(model_dict: dict) -> dict:
     # Add new assoc dict to new model dict
     new_model_dict['associations'] = new_assoc_list
 
-    # Reconstruct the attackers
-    if 'attackers' in model_dict:
-        attackers_info = model_dict['attackers']
-
     return new_model_dict
 
 
 def convert_model_dict_from_version_0_1(model_dict: dict) -> dict:
-    """
-    Convert model dict version 0.1 to 0.2
+    """Convert model dict version 0.1 to 0.2
 
     Arguments:
+    ---------
     model_dict  - the dictionary containing the serialized model
 
     Returns:
+    -------
     A dictionary containing the version 0.2 equivalent serialized model
-    """
 
+    """
     new_model_dict = {}
 
     # Meta data and assets format did not change from version 0.1
@@ -176,14 +168,14 @@ def convert_model_dict_from_version_0_1(model_dict: dict) -> dict:
     new_attackers_dict: dict[int, dict] = {}
     attackers_dict: dict = model_dict.get('attackers', {})
     for attacker_id, attacker_dict in attackers_dict.items():
-        attacker_id = int(attacker_id) # JSON compatibility
+        attacker_id = int(attacker_id)  # JSON compatibility
         new_attackers_dict[attacker_id] = {}
         new_attackers_dict[attacker_id]['name'] = attacker_dict['name']
         new_entry_points_dict = {}
 
         entry_points_dict = attacker_dict['entry_points']
         for asset_id, attack_steps in entry_points_dict.items():
-                asset_id = int(asset_id) # JSON compatibility
+                asset_id = int(asset_id)  # JSON compatibility
                 asset_name = new_assets_dict[asset_id]['name']
                 new_entry_points_dict[asset_name] = {
                     'asset_id': asset_id,
@@ -198,16 +190,17 @@ def convert_model_dict_from_version_0_1(model_dict: dict) -> dict:
 
 
 def convert_model_dict_from_version_0_2(model_dict: dict) -> dict:
-    """
-    Convert model dict version 0.2 to 0.3
+    """Convert model dict version 0.2 to 0.3
 
     Arguments:
+    ---------
     model_dict  - the dictionary containing the serialized model
 
     Returns:
+    -------
     A dictionary containing the version 0.3 equivalent serialized model
-    """
 
+    """
     new_model_dict = {}
 
     # Meta data and assets format did not change from version 0.1

maltoolbox/visualization/__init__.py

@@ -1,11 +1,11 @@
+from .draw_io_utils import create_drawio_file_with_images
 from .graphviz_utils import render_attack_graph, render_model
 from .neo4j_utils import ingest_attack_graph_neo4j, ingest_model_neo4j
-from .draw_io_utils import create_drawio_file_with_images
 
 __all__ = [
-    'render_attack_graph',
-    'render_model',
+    'create_drawio_file_with_images',
     'ingest_attack_graph_neo4j',
     'ingest_model_neo4j',
-    'create_drawio_file_with_images'
+    'render_attack_graph',
+    'render_model',
 ]